load distributed config

ai-hero · May 5, 2024 · be0d893 · be0d893
1 parent 7867029
commit be0d893
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/launch.py b/launch.py
@@ -9,8 +9,8 @@
 
 def train(training_config_file: str = "/mnt/config/training/config.yaml") -> None:
     """Run Training."""
-    training_config = TrainingJob.load(training_config_file, is_distributed=int(os.getenv("WORLD_SIZE", 1)) > 1)
-    TrainingJobRunner(training_config).run()
+    training_config = TrainingJob.load(training_config_file)
+    TrainingJobRunner(training_config, is_distributed=int(os.getenv("WORLD_SIZE", 1)) > 1).run()
 
 
 def infer(batch_inference_config_file: str = "/mnt/config/batch_inference/config.yaml") -> None: