Allow installing only tf or torch

rlworkgroup · Mar 11, 2021 · 8200b0c · 8200b0c
1 parent 62bbc5c
commit 8200b0c
Showing 1 changed file with 10 additions and 4 deletions.
diff --git a/setup.py b/setup.py
@@ -22,10 +22,6 @@
     'scikit-image',
     'scipy',
     'setproctitle>=1.0',
-    'tensorflow>=1.14',
-    'tensorflow-probability>=0.11.0',
-    'torch>=1.0.0,!=1.5.0',
-    'torchvision>=0.2.1',
 ]
 
 if sys.version_info < (3, 7):
@@ -34,6 +30,16 @@
 # Dependencies for optional features
 EXTRAS = {}
 
+EXTRAS['tensorflow'] = [
+    'tensorflow>=1.14',
+    'tensorflow-probability>=0.11.0',
+]
+
+EXTRAS['torch'] = [
+    'torch>=1.0.0,!=1.5.0',
+    'torchvision>=0.2.1',
+]
+
 EXTRAS['gym'] = [
     f'gym[atari,box2d,classic_control]=={GYM_VERSION}',
 ]