bug with saving

montrealrobotics · Oct 5, 2023 · 3821026 · 3821026
1 parent 50b8790
commit 3821026
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/cleanrl/ppo_continuous_action_wandb.py b/cleanrl/ppo_continuous_action_wandb.py
@@ -376,7 +376,7 @@ def train_risk(cfg, model, data, criterion, opt, device):
         opt.step()
 
         net_loss += loss.item()
-    torch.save(self.model.state_dict(), os.path.join(wandb.run.dir, "risk_model.pt"))
+    torch.save(model.state_dict(), os.path.join(wandb.run.dir, "risk_model.pt"))
     wandb.save("risk_model.pt")
     model.eval()
     return net_loss