fixed GraphData bug

CederGroupHub · Sep 15, 2023 · a6528ab · a6528ab
1 parent 1c745b3
commit a6528ab
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/chgnet/data/dataset.py b/chgnet/data/dataset.py
@@ -340,7 +340,7 @@ def __getitem__(self, idx) -> tuple[CrystalGraph, dict[str, Tensor]]:
                         force = self.labels[mp_id][graph_id][self.force_key]
                         targets["f"] = torch.tensor(force, dtype=datatype)
                     elif key == "s":
-                        stress = self.labels[mp_id][graph_id][self.magmom_key]
+                        stress = self.labels[mp_id][graph_id][self.stress_key]
                         # Convert VASP stress
                         targets["s"] = torch.tensor(stress, dtype=datatype) * (-0.1)
                     elif key == "m":