Update

[ghstack-poisoned]
pytorch · Dec 15, 2024 · 647dacc · 647dacc
1 parent f601cce
commit 647dacc
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/sota-implementations/ppo/ppo_atari.py b/sota-implementations/ppo/ppo_atari.py
@@ -236,6 +236,8 @@ def update(batch, num_network_updates):
                 with torch.no_grad(), timeit("adv"):
                     torch.compiler.cudagraph_mark_step_begin()
                     data = adv_module(data)
+                    if compile_mode:
+                        data = data.clone()
                 with timeit("rb - extend"):
                     # Update the data buffer
                     data_reshape = data.reshape(-1)

diff --git a/sota-implementations/ppo/ppo_mujoco.py b/sota-implementations/ppo/ppo_mujoco.py
@@ -225,6 +225,9 @@ def update(batch, num_network_updates):
                 with torch.no_grad(), timeit("adv"):
                     torch.compiler.cudagraph_mark_step_begin()
                     data = adv_module(data)
+                    if compile_mode:
+                        data = data.clone()
+
                 with timeit("rb - extend"):
                     # Update the data buffer
                     data_reshape = data.reshape(-1)