Merge branch 'main' into nkem/train-encoder

ZettaAI · Nov 28, 2023 · 8f0872a · 8f0872a
2 parents d8dda04 + 8b754d3
commit 8f0872a
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/zetta_utils/training/lightning/train.py b/zetta_utils/training/lightning/train.py
@@ -141,6 +141,7 @@ def lightning_train(
                 )
             train_args[k] = arg_spec
 
+    nproc_per_node = trainer.num_devices if nproc_per_node < 0 else nproc_per_node
     _lightning_train_remote(
         execution_id,
         cluster_info=cluster_info,