[dynamo][yolov3] Track UnspecializedNNModuleVariable for mutation (#128269)

Fixes https://github.com/pytorch/pytorch/issues/101168

Pull Request resolved: https://github.com/pytorch/pytorch/pull/128269
Approved by: https://github.com/jansel
ghstack dependencies: #128295, #126578, #128268, #128254

benchmarks/dynamo/ci_expected_accuracy/aot_eager_torchbench_inference.csv

@@ -378,4 +378,4 @@ vision_maskrcnn,pass,17
 
 
 
-yolov3,pass,2
+yolov3,pass,0

benchmarks/dynamo/ci_expected_accuracy/aot_eager_torchbench_training.csv

@@ -286,4 +286,4 @@ vision_maskrcnn,pass,34
 
 
 
-yolov3,pass,9
+yolov3,fail_accuracy,8

benchmarks/dynamo/ci_expected_accuracy/aot_inductor_torchbench_inference.csv

@@ -350,4 +350,4 @@ vision_maskrcnn,fail_to_run,0
 
 
 
-yolov3,fail_to_run,0
+yolov3,pass,0

benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_torchbench_freezing_inference.csv

@@ -338,4 +338,4 @@ vision_maskrcnn,pass,28
 
 
 
-yolov3,pass,2
+yolov3,pass,0

benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_torchbench_inference.csv

@@ -338,4 +338,4 @@ vision_maskrcnn,pass,28
 
 
 
-yolov3,pass,2
+yolov3,pass,0

benchmarks/dynamo/ci_expected_accuracy/cu124/aot_inductor_torchbench_inference.csv

@@ -350,4 +350,4 @@ vision_maskrcnn,fail_to_run,0
 
 
 
-yolov3,fail_to_run,0
+yolov3,pass,0

benchmarks/dynamo/ci_expected_accuracy/dynamic_aot_eager_torchbench_inference.csv

@@ -374,4 +374,4 @@ vision_maskrcnn,pass,17
 
 
 
-yolov3,pass,2
+yolov3,pass,0

benchmarks/dynamo/ci_expected_accuracy/dynamic_aot_eager_torchbench_training.csv

@@ -282,4 +282,4 @@ vision_maskrcnn,pass,34
 
 
 
-yolov3,pass,9
+yolov3,fail_accuracy,8

benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_inductor_torchbench_inference.csv

@@ -298,4 +298,4 @@ vision_maskrcnn,pass,28
 
 
 
-yolov3,pass,2
+yolov3,pass,0

benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_torchbench_inference.csv

@@ -374,4 +374,4 @@ vision_maskrcnn,pass,17
 
 
 
-yolov3,pass,2
+yolov3,pass,0

benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_torchbench_training.csv

@@ -282,4 +282,4 @@ vision_maskrcnn,pass,34
 
 
 
-yolov3,pass,9
+yolov3,pass,8

benchmarks/dynamo/ci_expected_accuracy/dynamo_eager_torchbench_inference.csv

@@ -378,4 +378,4 @@ vision_maskrcnn,pass,17
 
 
 
-yolov3,pass,2
+yolov3,pass,0

benchmarks/dynamo/ci_expected_accuracy/dynamo_eager_torchbench_training.csv

@@ -286,4 +286,4 @@ vision_maskrcnn,pass,34
 
 
 
-yolov3,pass,9
+yolov3,pass,8

benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_inference.csv

@@ -378,4 +378,4 @@ vision_maskrcnn,pass,17
 
 
 
-yolov3,pass,2
+yolov3,pass,0

benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_training.csv

@@ -286,4 +286,4 @@ vision_maskrcnn,pass,34
 
 
 
-yolov3,pass,9
+yolov3,pass,8

test/inductor/test_binary_folding.py

@@ -56,7 +56,7 @@ class BinaryFoldingTemplate(TestCase):
                     self.use_scalar = scalar
                     tensor_size = [1 for _ in range(self.conv.weight.ndim)]
                     tensor_size[1] = self.conv.weight.size(0)
-                    self.tensor = (
+                    self.tensor = torch.nn.Parameter(
                         add_tensor
                         if add_tensor is not None
                         else torch.rand(tensor_size).to(device)
@@ -136,7 +136,11 @@ class BinaryFoldingTemplate(TestCase):
                 nn.Conv2d,
                 pytorch_op,
                 False,
-                add_tensor=torch.rand(32, 1, 32).to(self.device),
+                add_tensor=torch.rand(
+                    32,
+                    1,
+                    32,
+                ).to(self.device),
                 expect_success=False,
             )
 
@@ -156,7 +160,7 @@ class BinaryFoldingTemplate(TestCase):
                 nn.Conv2d,
                 pytorch_op,
                 False,
-                add_tensor=torch.tensor([2]).to(torch.int).to(self.device),
+                add_tensor=torch.tensor([2]).to(torch.float64).to(self.device),
                 expect_success=False,
             )
 

torch/_dynamo/output_graph.py

@@ -752,7 +752,13 @@ class OutputGraph:
         **options,
     ):
         if is_dynamic_nn_module(target, self.root_tx.export):
-            return variables.UnspecializedNNModuleVariable(target, **options)
+            result = variables.UnspecializedNNModuleVariable(target, **options)
+            if not SideEffects.cls_supports_mutation_side_effects(type(target)):
+                # don't allow STORE_ATTR mutation with custom __setattr__
+                return result
+            return self.root_tx.output.side_effects.track_object_existing(
+                target, result
+            )
 
         options = dict(options)
         assert "source" in options