Hello,

Thanks for making this very useful package! Perhaps this is already in your radar but I wanted to flag the following deprecation warning in version 0.4.38 of jax

  /home/imendoza/.cache/pypoetry/virtualenvs/bliss-deblender-W0ocb7P7-py3.10/lib/python3.10/site-packages/torch_jax_interop/to_jax.py:107: DeprecationWarning: Calling from_dlpack with a DLPack tensor is deprecated. The argument to from_dlpack should be an array from another framework that implements the __dlpack__ protocol.
    return jax_from_dlpack(dlpack, copy=False)

not sure how difficult it might be to fix but would be great if this pacakge stayed compatible with jax going forward! Thanks!
(I would offer to help but I am really not an expert on any low level pytorch / jax programming)