diff --git a/setup.py b/setup.py
index 5527a12..409b1e7 100644
--- a/setup.py
+++ b/setup.py
@@ -13,8 +13,10 @@ setup(
         'torch',
         'numpy',
         'scipy',
-        'tensorboard',
     ],
+    extras_require={
+        "visualization":  ["tensorboard"],
+    },
     scripts=[
         'scripts/m2m.py',
     ]