Improve code in get_atomic_types (#469)

metatensor · Feb 1, 2025 · 71c8d49 · 71c8d49
1 parent 68bd9f5
commit 71c8d49
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/metatrain/utils/data/dataset.py b/src/metatrain/utils/data/dataset.py
@@ -190,8 +190,8 @@ def get_atomic_types(datasets: Union[Dataset, List[Dataset]]) -> List[int]:
 
     types = set()
     for dataset in datasets:
-        for index in range(len(dataset)):
-            system = dataset[index]["system"]
+        for sample in dataset:
+            system = sample["system"]
             types.update(set(system.types.tolist()))
 
     return sorted(types)