Update

[ghstack-poisoned]
pytorch · Dec 18, 2024 · ace76f3 · ace76f3
1 parent 595808e
commit ace76f3
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 9 deletions.
diff --git a/sota-implementations/ppo/ppo_atari.py b/sota-implementations/ppo/ppo_atari.py
@@ -245,10 +245,11 @@ def update(batch, num_network_updates):
                     data_buffer.extend(data_reshape)
 
                 for k, batch in enumerate(data_buffer):
-                    torch.compiler.cudagraph_mark_step_begin()
-                    loss, num_network_updates = update(
-                        batch, num_network_updates=num_network_updates
-                    )
+                    with timeit("update"):
+                        torch.compiler.cudagraph_mark_step_begin()
+                        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(

diff --git a/sota-implementations/ppo/ppo_mujoco.py b/sota-implementations/ppo/ppo_mujoco.py
@@ -239,11 +239,12 @@ def update(batch, num_network_updates):
                     data_buffer.extend(data_reshape)
 
                 for k, batch in enumerate(data_buffer):
-                    torch.compiler.cudagraph_mark_step_begin()
-                    loss, num_network_updates = update(
-                        batch, num_network_updates=num_network_updates
-                    )
-                    loss = loss.clone()
+                    with timeit("update"):
+                        torch.compiler.cudagraph_mark_step_begin()
+                        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"