Update

[ghstack-poisoned]
pytorch · Dec 15, 2024 · 9f08541 · 9f08541
2 parents f872d5c + aaec905
commit 9f08541
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/sota-implementations/ppo/ppo_atari.py b/sota-implementations/ppo/ppo_atari.py
@@ -181,7 +181,7 @@ def update(batch, num_network_updates):
 
         # Update the networks
         optim.step()
-        return loss.detach().set("alpha", alpha), num_network_updates.clone()
+        return loss.detach().set("alpha", alpha), num_network_updates
     if cfg.compile.compile:
         update = compile_with_warmup(update, mode=compile_mode, warmup=1)
         adv_module = compile_with_warmup(adv_module, mode=compile_mode, warmup=1)
@@ -246,6 +246,8 @@ def update(batch, num_network_updates):
                     loss, num_network_updates = update(
                         batch, num_network_updates=num_network_updates
                     )
+                    loss = loss.clone()
+                    num_network_updates = num_network_updates.clone()
                     losses[j, k] = loss.select(
                         "loss_critic", "loss_entropy", "loss_objective"
                     )

diff --git a/sota-implementations/ppo/ppo_mujoco.py b/sota-implementations/ppo/ppo_mujoco.py
@@ -166,7 +166,7 @@ def update(batch, num_network_updates):
 
         # Update the networks
         optim.step()
-        return loss.detach().set("alpha", alpha), num_network_updates.clone()
+        return loss.detach().set("alpha", alpha), num_network_updates
     if cfg.compile.compile:
         update = compile_with_warmup(update, mode=compile_mode, warmup=1)
         adv_module = compile_with_warmup(adv_module, mode=compile_mode, warmup=1)
@@ -235,6 +235,8 @@ def update(batch, num_network_updates):
                     loss, num_network_updates = update(
                         batch, num_network_updates=num_network_updates
                     )
+                    loss = loss.clone()
+                    num_network_updates = num_network_updates.clone()
                     losses[j, k] = loss.select(
                         "loss_critic", "loss_entropy", "loss_objective"
                     )