diff --git a/setup.py b/setup.py
index 9636a46..3656c05 100644
--- a/setup.py
+++ b/setup.py
@@ -15,8 +15,7 @@ setup(
         'torch>=1.2',
         'numpy',
         'scipy',
+        'matplotlib',
+        'tensorboard',
     ],
-    extras_require={
-        "visualization":  ["tensorboard"],
-    },
 )