fix breakage and add nested test

Dmytro Dzhulgakov · Dmytro Dzhulgakov · commit da0224dafaab · 2019-06-20T13:57:15.000-07:00
diff --git a/test/test_nn.py b/test/test_nn.py
@@ -4686,19 +4686,20 @@ def _load_from_state_dict(self, state_dict, prefix, local_metadata,
                 # skip some of the error handling
                 self.param.data.copy_(state_dict[prefix + "serialized"] - 1)
 
-        m = CustomState()
-        m.param[0] = 10
-        m.sub.weight[0, 0] = 555
+        # use sequential to verify nesting
+        m = nn.Sequential(CustomState())
+        m[0].param[0] = 10
+        m[0].sub.weight[0, 0] = 555
         state_dict = m.state_dict()
-        self.assertEqual(state_dict["serialized"].item(), 11)
-        self.assertIn("sub.weight", state_dict)
-        self.assertNotIn("param", state_dict)
+        self.assertEqual(state_dict["0.serialized"].item(), 11)
+        self.assertIn("0.sub.weight", state_dict)
+        self.assertNotIn("0.param", state_dict)
         del m
-        mm = CustomState()
-        self.assertEqual(mm.param[0].item(), 1)
+        mm = nn.Sequential(CustomState())
+        self.assertEqual(mm[0].param[0].item(), 1)
         mm.load_state_dict(state_dict)
-        self.assertEqual(mm.param[0].item(), 10)
-        self.assertEqual(mm.sub.weight[0, 0].item(), 555)
+        self.assertEqual(mm[0].param[0].item(), 10)
+        self.assertEqual(mm[0].sub.weight[0, 0].item(), 555)
 
     def test_parameter_assignment(self):
         l = nn.Linear(5, 5)
diff --git a/torch/jit/__init__.py b/torch/jit/__init__.py
@@ -1808,8 +1808,8 @@ def _get_methods(cls):
 _compiled_methods_whitelist = {
     'forward', 'register_buffer', 'register_parameter', 'add_module',
     '_apply', 'apply', 'cuda', 'cpu', 'to', 'type', 'float', 'double', 'half',
-    'state_dict', 'load_state_dict', '_load_from_state_dict',
-    '_named_members', 'parameters', 'named_parameters',
+    'state_dict', '_save_to_state_dict', 'load_state_dict',
+    '_load_from_state_dict', '_named_members', 'parameters', 'named_parameters',
     'buffers', 'named_buffers', 'children', 'named_children', 'modules',
     'named_modules', 'zero_grad', 'share_memory', '_get_name', 'extra_repr',
     '_slow_forward', '_tracing_name', 'eval', 'train',
