fix bug

allenai · Jul 12, 2024 · 9aedd25 · 9aedd25
1 parent 1a43d36
commit 9aedd25
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/allenact/utils/experiment_utils.py b/allenact/utils/experiment_utils.py
@@ -1208,9 +1208,10 @@ def download_checkpoint_from_wandb(checkpoint_path_dir_or_pattern, all_ckpt_dir,
         return ckpts_paths
     else:
         assert len(ckpt_steps) == 1
-        ckpt_fn = "{}-step-{}:latest".format(run_token, steps)
+        step = ckpt_steps[0]
+        ckpt_fn = "{}-step-{}:latest".format(run_token, step)
         artifact = api.artifact(ckpt_fn)
         _ = artifact.download("/tmp")
-        ckpt_dir = "{}/ckpt-{}.pt".format(all_ckpt_dir, steps)
+        ckpt_dir = "{}/ckpt-{}.pt".format(all_ckpt_dir, step)
         shutil.move("/tmp/ckpt.pt", ckpt_dir)
         return ckpt_dir