Fix predict_as_dataframe: single-pass, shuffle-safe, multi-dim

examples/04_training/01_train_dynedge.py

@@ -178,6 +178,10 @@ def main(
         gpus=config["fit"]["gpus"],
     )
 
+    if results is None:
+        # Only the global-zero rank holds the output in multi-device runs.
+        return
+
     # Save predictions and model to file
     db_name = path.split("/")[-1].split(".")[0]
     path = os.path.join(archive, db_name, run_name)

examples/04_training/02_train_tito_model.py

@@ -177,6 +177,10 @@ def main(
         gpus=config["fit"]["gpus"],
     )
 
+    if results is None:
+        # Only the global-zero rank holds the output in multi-device runs.
+        return
+
     # Save predictions and model to file
     db_name = path.split("/")[-1].split(".")[0]
     path = os.path.join(archive, db_name, run_name)

examples/04_training/03_train_dynedge_from_config.py

@@ -113,6 +113,10 @@ def main(
         additional_attributes=additional_attributes + ["event_no"],
         gpus=config.fit["gpus"],
     )
+
+    if results is None:
+        # Only the global-zero rank holds the output in multi-device runs.
+        return
     results.to_csv(f"{path}/results.csv")
 
 

examples/04_training/04_train_multiclassifier_from_configs.py

@@ -139,6 +139,10 @@ def main(
         additional_attributes=additional_attributes + ["event_no"],
         gpus=config.fit["gpus"],
     )
+
+    if results is None:
+        # Only the global-zero rank holds the output in multi-device runs.
+        return
     results.to_csv(f"{path}/results.csv")
 
 

examples/04_training/05_train_RNN_TITO.py

@@ -199,6 +199,10 @@ def main(
         gpus=config["fit"]["gpus"],
     )
 
+    if results is None:
+        # Only the global-zero rank holds the output in multi-device runs.
+        return
+
     # Save predictions and model to file
     db_name = path.split("/")[-1].split(".")[0]
     path = os.path.join(archive, db_name, run_name)

examples/04_training/06_train_icemix_model.py

@@ -205,6 +205,10 @@ def main(
         gpus=config["fit"]["gpus"],
     )
 
+    if results is None:
+        # Only the global-zero rank holds the output in multi-device runs.
+        return
+
     # Save predictions and model to file
     db_name = path.split("/")[-1].split(".")[0]
     path = os.path.join(archive, db_name, run_name)

examples/04_training/07_train_normalizing_flow.py

@@ -148,6 +148,10 @@ def main(
         gpus=config["fit"]["gpus"],
     )
 
+    if results is None:
+        # Only the global-zero rank holds the output in multi-device runs.
+        return
+
     # Save predictions and model to file
     db_name = path.split("/")[-1].split(".")[0]
     path = os.path.join(archive, db_name, run_name)

examples/04_training/08_train_grit_model.py

@@ -166,6 +166,10 @@ def main(
         gpus=config["fit"]["gpus"],
     )
 
+    if results is None:
+        # Only the global-zero rank holds the output in multi-device runs.
+        return
+
     # Save predictions and model to file
     db_name = path.split("/")[-1].split(".")[0]
     path = os.path.join(archive, db_name, run_name)