Install torch cpu version for tests

mseitzer · Mar 16, 2024 · 6211a10 · 6211a10
1 parent 88d0de4
commit 6211a10
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/noxfile.py b/noxfile.py
@@ -15,6 +15,11 @@ def lint(session):
 
 @nox.session(python=["3.8", "3.9", "3.10", "3.11", "3.12"])
 def tests(session):
+    session.install(
+        'torch==2.2.1', 
+        'torchvision', 
+        '--index-url', 'https://download.pytorch.org/whl/cpu'
+    )
     session.install('.')
     session.install('pytest')
     session.install('pytest-mock')