diff --git a/setup.py b/setup.py
index 0877f90..e635fec 100644
--- a/setup.py
+++ b/setup.py
@@ -24,7 +24,9 @@ setup(
     'pytest-runner',
   ],
   tests_require=[
-    'pytest'
+    'pytest',
+    'torch==1.12.1',
+    'torchvision==0.13.1'
   ],
   classifiers=[
     'Development Status :: 4 - Beta',