fix bug: use provided revision instead of hardcoded one

This commit is contained in:
Alexander Soare 2024-03-27 13:08:47 +00:00
parent 64c8851c40
commit b6353964ba
1 changed files with 1 additions and 1 deletions

View File

@ -268,7 +268,7 @@ if __name__ == "__main__":
# TODO(alexander-soare): Save and load stats in trained model directory.
stats_path = None
elif args.hub_id is not None:
folder = Path(snapshot_download(args.hub_id, revision="v1.0"))
folder = Path(snapshot_download(args.hub_id, revision=args.revision))
cfg = hydra.initialize(config_path=str(_relative_path_between(folder, Path(__file__).parent)))
cfg = hydra.compose("config", args.overrides)
cfg.policy.pretrained_model_path = folder / "model.pt"