allow wandb to upload ckpts

allenai · Jul 5, 2024 · 04c5b80 · 04c5b80
1 parent bf93d8d
commit 04c5b80
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/allenact/algorithms/onpolicy_sync/runner.py b/allenact/algorithms/onpolicy_sync/runner.py
@@ -1109,9 +1109,16 @@ def update_keys_metric(
                             storage_uuid_to_total_experiences_key[storage_uuid]
                         )
 
-        assert all_equal(
-            checkpoint_file_name
-        ), f"All {mode} logging packages must have the same checkpoint_file_name."
+        if any(checkpoint_file_name):
+            ckpt_to_store = None
+            for ckpt in checkpoint_file_name:
+                if ckpt is not None:
+                    ckpt_to_store = ckpt
+            assert ckpt_to_store is not None
+            checkpoint_file_name = [ckpt_to_store]
+        # assert all_equal(
+        #     checkpoint_file_name
+        # ), f"All {mode} logging packages must have the same checkpoint_file_name."
 
         message = [
             f"{mode.upper()}: {training_steps} rollout steps ({pkgs[0].storage_uuid_to_total_experiences})"