dont emit non mutable weights

Differential Revision: D61888453

Pull Request resolved: #4938

Author: JacobSzwejbka

exir/passes/weights_to_outputs_pass.py

@@ -53,11 +53,13 @@ def weights_to_outputs_pass(
             break
     assert output_node is not None
 
-    # Get place holder nodes with gradients
+    # Get input nodes that are weights with an associated gradient
     placeholder_nodes = [
         node
         for node in gm.graph.nodes
-        if node.op == "placeholder" and node.target in inputs_to_params.keys()
+        if node.op == "placeholder"
+        and node.target in inputs_to_params.keys()
+        and inputs_to_params[node.target] in grad_targets
     ]
 
     # Flag these placeholder nodes as having a gradient attached so that memory planning will operate on them.

exir/tests/test_joint_graph.py

@@ -26,10 +26,13 @@ class Module(torch.nn.Module):
             def __init__(self):
                 super().__init__()
                 self.linear = torch.nn.Linear(3, 3)
+                self.linear_no_train = torch.nn.Linear(3, 3)
+                for param in self.linear_no_train.parameters():
+                    param.requires_grad = False
                 self.loss = torch.nn.CrossEntropyLoss()
 
             def forward(self, x, y):
-                return self.loss(self.linear(x).softmax(dim=0), y)
+                return self.loss(self.linear_no_train(self.linear(x)).softmax(dim=0), y)
 
         m = Module()
         example_inputs = (torch.ones(3), torch.tensor([1.0, 0.0, 0.0]))