Update on "[quant][pt2] Preserve source_fn_stack after QAT fusion"

andrewor14 · andrewor14 · commit 3341aef799f9 · 2023-10-10T11:03:01.000-07:00
Test Plan: python test/test_quantization.py TestQuantizePT2EQAT.test_qat_preserve_source_fn_stack Reviewers: jerryzh168, kimishpatel Subscribers: jerryzh168, kimishpatel, supriyar Differential Revision: [D50101253](https://our.internmc.facebook.com/intern/diff/D50101253) [ghstack-poisoned]
diff --git a/test/quantization/pt2e/test_quantize_pt2e_qat.py b/test/quantization/pt2e/test_quantize_pt2e_qat.py
@@ -532,6 +532,7 @@ def test_qat_preserve_source_fn_stack(self):
         """
         Test whether `source_fn_stack` is preserved after QAT fusion.
         """
+
         class M(torch.nn.Module):
             def __init__(self):
                 super().__init__()
@@ -579,13 +580,15 @@ def get_conv_weight_and_bias(conv_node: torch.fx.Node):
             assert isinstance(weight_node, torch.fx.Node)
             assert isinstance(bias_node, torch.fx.Node)
             return (weight_node, bias_node)
+
         first_conv_weight, first_conv_bias = get_conv_weight_and_bias(first_conv)
         second_conv_weight, second_conv_bias = get_conv_weight_and_bias(second_conv)
 
         # Assert that each set of conv, conv weight, and conv bias are in the same partition
         def get_source_fn(node: torch.fx.Node):
             # E.g. [('l__self___backbone1_conv', <class 'torch.nn.modules.conv.Conv2d'>)]
             return node.meta["source_fn_stack"][0][0]
+
         self.assertEqual(get_source_fn(first_conv), get_source_fn(first_conv_weight))
         self.assertEqual(get_source_fn(first_conv), get_source_fn(first_conv_bias))
         self.assertEqual(get_source_fn(second_conv), get_source_fn(second_conv_weight))
