Even more mock imports

MKorablyov · Jan 30, 2020 · 0cba288 · 0cba288
1 parent 4880fc6
commit 0cba288
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/docs/conf.py b/docs/conf.py
@@ -33,7 +33,15 @@ class Mock(MagicMock):
     def __getattr__(cls, name):
         return MagicMock()
 
-MOCK_MODULES = ['mpi4py', 'torch', 'torch.optim', 'torch.nn']
+MOCK_MODULES = ['mpi4py', 
+                'torch', 
+                'torch.optim', 
+                'torch.nn',
+                'torch.distributions',
+                'torch.distributions.normal',
+                'torch.distributions.categorical',
+                'torch.nn.functional',
+                ]
 sys.modules.update((mod_name, Mock()) for mod_name in MOCK_MODULES)
 
 # Finish imports