Update

[ghstack-poisoned]
pytorch · Dec 15, 2024 · f601cce · f601cce
1 parent f3fd558
commit f601cce
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
@@ -234,6 +234,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
@@ -223,7 +223,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)