Fix bug with source map

pytorch · Jun 3, 2023 · 7fbd56d · 7fbd56d
1 parent 9c3b0d4
commit 7fbd56d
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/torch/_dynamo/variables/optimizer.py b/torch/_dynamo/variables/optimizer.py
@@ -45,8 +45,9 @@ def map_grads_to_sources(self):
             group_source = GetItemSource(AttrSource(self.source, "param_groups"), g_ind)
             for p_ind, p in enumerate(group["params"]):
                 if p.grad is not None:
-                    self.grad_to_source[p.grad] = GetItemSource(
-                        GetItemSource(group_source, "params"), p_ind
+                    self.grad_to_source[p.grad] = AttrSource(
+                        GetItemSource(GetItemSource(group_source, "params"), p_ind),
+                        "grad",
                     )
 
     def var_getattr(self, tx, name):