Add name of struct

GreeleyGroup · Sep 20, 2023 · b3e9246 · b3e9246
1 parent 003f669
commit b3e9246
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/data.py b/src/data.py
@@ -109,6 +109,9 @@ def process(self, z_cutoffs, node_features, edge_features, max_atoms=12,
                 )
                 data_objects.append(data_obj)
 
+            # Add name of structure
+            data_objects.append(name)
+
             # Save data objects
             torch.save(data_objects, self.processed_path / f"data_{i}.pt")