pass entire config to make_optimizer

2024-09-02 08:20:17 +00:00
parent 3034272229
commit 06fc9b89e1
7 changed files with 24 additions and 29 deletions
--- a/tests/scripts/save_policy_to_safetensors.py
+++ b/tests/scripts/save_policy_to_safetensors.py
@@ -39,7 +39,7 @@ def get_policy_stats(env_name, policy_name, extra_overrides):
    dataset = make_dataset(cfg)
    policy = make_policy(cfg, dataset_stats=dataset.stats)
    policy.train()
-    optimizer, _ = policy.make_optimizer_and_scheduler(**cfg.training)
+    optimizer, _ = policy.make_optimizer_and_scheduler(cfg)

    dataloader = torch.utils.data.DataLoader(
        dataset,