adding docstring and from_pretrained/save_pretrained

2024-04-02 22:45:21 +02:00
parent 11cbf1bea1
commit 4751642ace
9 changed files with 79 additions and 30 deletions
--- a/examples/3_train_policy.py
+++ b/examples/3_train_policy.py
@@ -50,6 +50,6 @@ for offline_step in trange(cfg.offline_steps):
        print(train_info)

 # Save the policy, configuration, and normalization stats for later use.
-policy.save(output_directory / "model.pt")
+policy.save_pretrained(output_directory / "model.pt")
 OmegaConf.save(cfg, output_directory / "config.yaml")
 torch.save(offline_buffer.transform[-1].stats, output_directory / "stats.pth")