import torch

lucidrains · lucidrains · commit 60ca207fe1d7 · 2021-10-25T18:18:47.000-07:00
diff --git a/jax2torch/jax2torch.py b/jax2torch/jax2torch.py
@@ -1,5 +1,6 @@
 # https://gist.github.com/mattjj/e8b51074fed081d765d2f3ff90edf0e9
 
+import torch
 from jax import dlpack as jax_dlpack
 from torch.utils import dlpack as torch_dlpack
 
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'jax2torch',
   packages = find_packages(exclude=[]),
-  version = '0.0.2',
+  version = '0.0.3',
   license='MIT',
   description = 'Jax 2 Torch',
   author = 'Phil Wang',
