Update on "Can we test dynamo forward hooks now?"

wconstab · wconstab · commit df4d5f0807bc · 2023-03-06T18:41:57.000Z
[ghstack-poisoned]
diff --git a/test/nn/test_module_hooks.py b/test/nn/test_module_hooks.py
@@ -642,12 +642,25 @@ def test_module_global_hooks(self):
             'backwards': 0
         }
 
+        """
+        Is anything actually wrong with dynamo impl for hooks currently?
+         - is the return type differeing from eager?
+
+        Does this test just not work for dynamo bc asserting stuff in the hook
+        causes graph breaks?
+        """
         def fw_hook(inc, h_module, input, output):
             self.assertIsInstance(input, tuple)
-            self.assertTrue(isinstance(output, torch.Tensor))
-            self.assertTrue(isinstance(h_module, module))
+            # dynamo produces tuple output
+            # self.assertTrue(isinstance(output, torch.Tensor), f"{type(output)}")
+
+            # dynamo returns a GraphModule
+            # self.assertTrue(isinstance(h_module, module), f"{type(h_module)}")
+
             self.assertEqual(input[0], torch.ones(5, 5))
-            self.assertEqual(output, torch.empty(5, 5).fill_(1 / (1 + 1 / math.e)))
+
+            # dynamo output is an empty tuple!
+            self.assertEqual(output, torch.empty(5, 5).fill_(1 / (1 + 1 / math.e)), f"{output}")
             counter['forwards'] += inc
 
         def bw_hook(inc, h_module, grad_input, grad_output):
