Update

[ghstack-poisoned]
pytorch · Dec 15, 2024 · 2f29fa9 · 2f29fa9
2 parents b24ad47 + 1b4adfd
commit 2f29fa9
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/sota-implementations/ppo/ppo_atari.py b/sota-implementations/ppo/ppo_atari.py
@@ -235,6 +235,7 @@ def update(batch, num_network_updates):
 
                 # Compute GAE
                 with torch.no_grad(), timeit("adv"):
+                    torch.compiler.cudagraph_mark_step_begin()
                     data = adv_module(data)
                 with timeit("rb - extend"):
                     # Update the data buffer

diff --git a/sota-implementations/ppo/ppo_mujoco.py b/sota-implementations/ppo/ppo_mujoco.py
@@ -224,7 +224,8 @@ def update(batch, num_network_updates):
 
                 # Compute GAE
                 with torch.no_grad(), timeit("adv"):
-                    data = adv_module(data.to(device))
+                    torch.compiler.cudagraph_mark_step_begin()
+                    data = adv_module(data)
                 with timeit("rb - extend"):
                     # Update the data buffer
                     data_reshape = data.reshape(-1)