Merge branch 'datafix' of https://github.com/CaRoLZhangxy/deepmd-kit …

…into datafix
deepmodeling · Apr 1, 2024 · 722bac6 · 722bac6
2 parents 7ce216e + 0604d16
commit 722bac6
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
@@ -1053,7 +1053,9 @@ def get_data(self, is_train=True, task_key="Default"):
                 except StopIteration:
                     # Refresh the status of the dataloader to start from a new epoch
                     with torch.device("cpu"):
-                        self.training_data[task_key] = self.training_dataloader[task_key]
+                        self.training_data[task_key] = self.training_dataloader[
+                            task_key
+                        ]
                     batch_data = next(iter(self.training_data[task_key]))
             else:
                 if self.validation_data[task_key] is None: