add submodules to sys.modules so their attributes can be pickled (#53107)

mattip · facebook-github-bot · commit fbf9745c85b6 · 2021-03-02T10:47:21.000-08:00
Summary: Fixes #38137 As mentioned in the issue, this is a workaround for [python issue 43367](https://bugs.python.org/issue43367). There are a number of other places where `sys.modules` is modified, if something changes in python perhaps those should be reviewed as well. Pull Request resolved: #53107 Reviewed By: zou3519 Differential Revision: D26753571 Pulled By: ezyang fbshipit-source-id: 2bda03bab39ff9ca58ce4bc13befe021da91b9c4
diff --git a/test/test_nn.py b/test/test_nn.py
@@ -14926,6 +14926,14 @@ def test_invalid_functions(self):
         with self.assertRaisesRegex(ValueError, 'uninitialized parameter'):
             param + param
 
+class TestFunctionalPickle(TestCase):
+
+    # issue gh-38137
+    def test_pickle_softsign(self):
+        # Make sure it does not throw an exception
+        s = pickle.dumps(F.softsign)
+    
+
 instantiate_device_type_tests(TestNNDeviceType, globals())
 
 if __name__ == '__main__':
diff --git a/torch/__init__.py b/torch/__init__.py
@@ -230,6 +230,19 @@ def _load_global_deps():
             if name[0] != '_' and
             not name.endswith('Base')]
 
+if not TYPE_CHECKING:
+    # issue 38137 and python issue 43367. Submodules of a C extension are
+    # non-standard, and attributes of those submodules cannot be pickled since
+    # pickle expect to be able to import them as "from _C.sub import attr"
+    # which fails with "_C is not a package
+    for attr in dir(_C):
+        candidate = getattr(_C, attr)
+        if type(candidate) is type(_C):
+            # submodule
+            if f'torch._C.{attr}' not in sys.modules:
+                sys.modules[f'torch._C.{attr}'] = candidate
+ 
+
 ################################################################################
 # Define basic utilities
 ################################################################################
