move download_checkpoint_from_wandb from engine to runner

allenai · Jul 12, 2024 · ecc24a9 · ecc24a9
1 parent 25cbe11
commit ecc24a9
Showing 1 changed file with 9 additions and 8 deletions.
diff --git a/allenact/algorithms/onpolicy_sync/runner.py b/allenact/algorithms/onpolicy_sync/runner.py
@@ -534,14 +534,15 @@ def start_train(
 
  worker_ids = self.local_worker_ids(TRAIN_MODE_STR)
 
- if "wandb://" == checkpoint[:8]:
- ckpt_dir = "/tmp/wandb_ckpts"
- os.makedirs(ckpt_dir, exist_ok=True)
- checkpoint = download_checkpoint_from_wandb(
- checkpoint,
- ckpt_dir,
- only_allow_one_ckpt=True
- )
+ if checkpoint is not None:
+ if checkpoint[:8] == "wandb://":
+ ckpt_dir = "/tmp/wandb_ckpts"
+ os.makedirs(ckpt_dir, exist_ok=True)
+ checkpoint = download_checkpoint_from_wandb(
+ checkpoint,
+ ckpt_dir,
+ only_allow_one_ckpt=True
+ )
 
  model_hash = None
  for trainer_id in worker_ids: