How to lazy load headless ui Dialog in React?

I’m trying to lazy load a headless ui Dialog component but I haven’t been able to do it while also keeping transitions.

Here’s as close as I’ve gone:

// Modal.js

const Modal = ({
  isOpen
}) => {
  return (
    <Transition show={isOpen}>
      <Dialog>
        ...
      </Dialog> 
    </Transition>
  );
};

export default Modal;

const Modal = lazy(() => import("~/components/Modal"));

<Suspense>
  {isOpen ? <Modal isOpen={isOpen} ... />
</Suspense>

This works for the most part but I lose the transition when the component is unmounted since isOpen is false then it gets unmounted immediately. But If I remove it then the component isn’t lazy loaded when I open it but instead it loads on page load.

Is there a way to lazy load it while keeping the Transitions?