Add setup

NKI-AI · Jan 3, 2025 · bc7605e · bc7605e
1 parent b22a92a
commit bc7605e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -59,12 +59,13 @@ def finalize_options(self):
         "h5py==3.11.0",
         "omegaconf==2.3.0",
         "torch>=2.2.0",
-        "torchvision==0.18.0",
+        "torchvision",
         "scikit-image>=0.19.0",
         "scikit-learn>=1.0.1",
         "tensorboard>=2.7.0",
         "tqdm",
         "protobuf==3.20.2",
+        "einops",
     ],
     extras_require={
         "dev": [