Merge pull request #291 from Czaki/setup_fix

Incorporate czaki contribution

Author: FynnBe

setup.py

@@ -38,7 +38,7 @@
     extras_require={
         "test": ["pytest", "black", "mypy"],
         "dev": ["pre-commit"],
-        "pytorch": ["pytorch>=1.6", "torchvision", "cudatoolkit>=10.1"],
+        "pytorch": ["torch>=1.6", "torchvision"],
         "tensorflow": ["tensorflow"],
         "onnx": ["onnxruntime"],
     },