fix logging in multigpu training

molmod · Jul 20, 2024 · 95e5a82 · 95e5a82
1 parent b348ee6
commit 95e5a82
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/examples/water_train_validate.py b/examples/water_train_validate.py
@@ -14,6 +14,8 @@ def main():
         correlation=3,
         max_L=1,
         num_channels=16,
+        max_num_epochs=20,
+        swa=False,
     )
 
     train, valid = data.split(0.9, shuffle=True)

diff --git a/psiflow/models/mace_utils.py b/psiflow/models/mace_utils.py
@@ -77,7 +77,7 @@ def run(rank: int, args: argparse.Namespace, world_size: int) -> None:
 
     # Setup
     tools.set_seeds(args.seed)
-    tools.setup_logger(level=args.log_level, tag=tag, directory=args.log_dir)
+    tools.setup_logger(level=args.log_level, tag=tag, directory=args.log_dir, rank=rank)
     try:
         logging.info(f"MACE version: {mace.__version__}")
     except AttributeError: