From 8200b0c413e9c32aa6f559ac37d283064bac5bf2 Mon Sep 17 00:00:00 2001
From: Gitanshu Sardana <gitanshusardana@gmail.com>
Date: Tue, 27 Oct 2020 19:13:56 -0700
Subject: [PATCH] Allow installing only tf or torch

---
 setup.py | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/setup.py b/setup.py
index b4f394a69a..94135f88de 100644
--- 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}',
 ]