[quant] follow up fixes for prepare_fx/prepare_qat_fx calls in classyvision (#105) (pytorch#78660)

Summary:
X-link: https://github.com/fairinternal/ClassyVision/pull/105

As follow up for pytorch#76496, we fixes the TODOs in quantization tests
by providing correct example_inputs in the tests

Test Plan:
classyvision sandcastle and ossci

**Static Docs Preview: classyvision**
|[Full Site](https://our.intern.facebook.com/intern/staticdocs/eph/D36818665/V1/classyvision/)|

|**Modified Pages**|

Differential Revision: D36818665

Pull Request resolved: pytorch#78660
Approved by: https://github.com/vkuzo

Author: jerryzh168

torch/ao/quantization/utils.py

@@ -495,35 +495,29 @@ def forward(self, x, key1=3, key2=3):
     root = model
     fqn_to_example_inputs = {}
 
-    class InterceptionModule(type(model)):  # type: ignore[misc]
-        def __call__(self, *args, **kwargs):
-            orig_module_call = torch.nn.Module.__call__
-
-            def _patched_module_call(self, *args, **kwargs):
-                submodule_example_inputs = list(args).copy()
-                normalized_kwargs = _normalize_kwargs(self.forward, kwargs)
-                # minus 1 to skipping counting `self`
-                num_args = _get_num_pos_args(self.forward) - 1
-                num_to_pop = num_args - len(submodule_example_inputs)
-                while num_to_pop and normalized_kwargs:
-                    normalized_kwargs.popitem(last=False)
-                    num_to_pop -= 1
-                submodule_example_inputs.extend(normalized_kwargs.values())
-                submodule_example_inputs_tuple = tuple(submodule_example_inputs)
-                fqn = _get_path_of_module(root, self)
-                if fqn is not None:
-                    fqn_to_example_inputs[fqn] = submodule_example_inputs_tuple
-                return orig_module_call(self, *args, **kwargs)
-
-            torch.nn.Module.__call__ = _patched_module_call
-            super().__call__(*args, **kwargs)
-            torch.nn.Module.__call__ = orig_module_call
-
-    original_class = model.__class__
-    model.__class__ = InterceptionModule
-    model(*example_inputs)
-    model.__class__ = original_class
-
+    def _patched_module_call(self, *args, **kwargs):
+        submodule_example_inputs = list(args).copy()
+        normalized_kwargs = _normalize_kwargs(self.forward, kwargs)
+        # minus 1 to skipping counting `self`
+        num_args = _get_num_pos_args(self.forward) - 1
+        num_to_pop = num_args - len(submodule_example_inputs)
+        while num_to_pop and normalized_kwargs:
+            normalized_kwargs.popitem(last=False)
+            num_to_pop -= 1
+        submodule_example_inputs.extend(normalized_kwargs.values())
+        submodule_example_inputs_tuple = tuple(submodule_example_inputs)
+        fqn = _get_path_of_module(root, self)
+        if fqn is not None:
+            fqn_to_example_inputs[fqn] = submodule_example_inputs_tuple
+        return orig_module_call(self, *args, **kwargs)
+
+    orig_module_call = torch.nn.Module.__call__
+    torch.nn.Module.__call__ = _patched_module_call
+    try:
+        model(*example_inputs)
+    finally:
+        # restore the module call even if there is an exception
+        torch.nn.Module.__call__ = orig_module_call
     return fqn_to_example_inputs