Remove tensor name in collate_fn (#460)

metatensor · Jan 29, 2025 · 52d1eb7 · 52d1eb7
1 parent d358315
commit 52d1eb7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/metatrain/utils/data/dataset.py b/src/metatrain/utils/data/dataset.py
@@ -228,7 +228,7 @@ def collate_fn(batch: List[Dict[str, Any]]) -> Tuple[List, Dict[str, TensorMap]]
     targets.
     """
 
-    collated_targets = group_and_join(batch)
+    collated_targets = group_and_join(batch, join_kwargs={"remove_tensor_name": True})
     collated_targets = collated_targets._asdict()
     systems = collated_targets.pop("system")
     return systems, collated_targets