fix: correcting graph get retrieval

IntelLabs · May 24, 2024 · 2f8894c · 2f8894c
1 parent 4e87c5d
commit 2f8894c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/matsciml/models/base.py b/matsciml/models/base.py
@@ -2365,7 +2365,7 @@ def _toggle_input_grads(
                     for key in input_keys:
                         # set require grad for both point cloud and graph tensors
                         if "graph" in data:
-                            g = data.get("g")
+                            g = data.get("graph")
                             if isinstance(g, dgl.DGLGraph):
                                 if key in g.ndata:
                                     data["graph"].ndata[key].requires_grad_(True)