From 0cba2886047b7de82c2cad4321df5875db644d61 Mon Sep 17 00:00:00 2001
From: jachiam <jachiam@berkeley.edu>
Date: Thu, 30 Jan 2020 08:38:29 -0800
Subject: [PATCH] Even more mock imports

---
 docs/conf.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/docs/conf.py b/docs/conf.py
index aba5f0845..9a524d3b8 100644
--- 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