Error: Call to undefined method Drupal\commerce_checkout\Plugin\Commerce\CheckoutFlow\MultistepDefault::getFormDisplay()

Problem/Motivation

Does not work with Commerce order type fields.

Steps to reproduce

1. Enable this module (Field States UI)
2. Enable Commerce Checkout Order Fields →
3. Go to /admin/commerce/config/order-types/default/edit/fields and create couple fields, for example, simple text field and a boolean.
4. Go to /admin/commerce/config/order-types/default/edit/form-display and configure the text field to make it Visible on Target: field_boolean with Comparison type: checked.
5. Place a product to cart and go to checkout page to see:

The website encountered an unexpected error. Please try again later.
Error: Call to undefined method Drupal\commerce_checkout\Plugin\Commerce\CheckoutFlow\MultistepDefault::getFormDisplay() in Drupal\field_states_ui\FieldStateBase->applyState() (line 72 of modules/contrib/field_states_ui/src/FieldStateBase.php).

Drupal\field_states_ui\FieldStateBase->applyState() (Line: 308)
Drupal\field_states_ui\FieldStateManager->getStates() (Line: 265)
Drupal\field_states_ui\FieldStateManager->apply() (Line: 37)
field_states_ui_field_widget_multivalue_form_alter() (Line: 539)
Drupal\Core\Extension\ModuleHandler->alter() (Line: 549)
Drupal\Core\Extension\ModuleHandler->alterDeprecated() (Line: 122)
Drupal\Core\Field\WidgetBase->form() (Line: 181)
Drupal\Core\Entity\Entity\EntityFormDisplay->buildForm() (Line: 104)
Drupal\commerce_checkout_order_fields\Plugin\Commerce\CheckoutPane\OrderFields->buildPaneForm() (Line: 546)
Drupal\commerce_checkout\Plugin\Commerce\CheckoutFlow\CheckoutFlowWithPanesBase->buildForm()
call_user_func_array() (Line: 531)
Drupal\Core\Form\FormBuilder->retrieveForm() (Line: 278)
Drupal\Core\Form\FormBuilder->buildForm() (Line: 219)
Drupal\Core\Form\FormBuilder->getForm() (Line: 143)
Drupal\commerce_checkout\Controller\CheckoutController->formPage()
call_user_func_array() (Line: 123)
Drupal\Core\EventSubscriber\EarlyRenderingControllerWrapperSubscriber->Drupal\Core\EventSubscriber\{closure}() (Line: 564)
Drupal\Core\Render\Renderer->executeInRenderContext() (Line: 124)
Drupal\Core\EventSubscriber\EarlyRenderingControllerWrapperSubscriber->wrapControllerExecutionInRenderContext() (Line: 97)
Drupal\Core\EventSubscriber\EarlyRenderingControllerWrapperSubscriber->Drupal\Core\EventSubscriber\{closure}() (Line: 158)
Symfony\Component\HttpKernel\HttpKernel->handleRaw() (Line: 80)
Symfony\Component\HttpKernel\HttpKernel->handle() (Line: 58)
Drupal\Core\StackMiddleware\Session->handle() (Line: 48)
Drupal\Core\StackMiddleware\KernelPreHandle->handle() (Line: 106)
Drupal\page_cache\StackMiddleware\PageCache->pass() (Line: 85)
Drupal\page_cache\StackMiddleware\PageCache->handle() (Line: 50)
Drupal\ban\BanMiddleware->handle() (Line: 48)
Drupal\Core\StackMiddleware\ReverseProxyMiddleware->handle() (Line: 51)
Drupal\Core\StackMiddleware\NegotiationMiddleware->handle() (Line: 23)
Stack\StackedHttpKernel->handle() (Line: 708)
Drupal\Core\DrupalKernel->handle() (Line: 19)

Tested with regular content type fields and it works perfectly. Would be nice if this module also worked with Drupal Commerce order type fields.