correct torch package name

cpnota · Jan 14, 2020 · ea91901 · ea91901
1 parent eebef5d
commit ea91901
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -24,7 +24,8 @@
     ],
     extras_require={
         "pytorch": [
-            "pytorch",
+            "torch",
+            "torchvision",
             "tensorboard"
         ],
         "docs": [