Update

[ghstack-poisoned]
pytorch · vmoens · Dec 16, 2024 · Dec 14, 2024 · Dec 14, 2024 · Dec 14, 2024
commit 65bb4f3ac93a7ca042df60820bb963bb29f2fbbb
diff --git a/sota-implementations/iql/discrete_iql.py b/sota-implementations/iql/discrete_iql.py
@@ -184,6 +184,7 @@ def update(sampled_tensordict):
                         sampled_tensordict = replay_buffer.sample().to(device)
 
                     with timeit("training - update"):
+                        torch.compiler.cudagraph_mark_step_begin()
                         metadata = update(sampled_tensordict)
                     # update priority
                     if prb:

diff --git a/sota-implementations/iql/iql_offline.py b/sota-implementations/iql/iql_offline.py
@@ -139,6 +139,7 @@ def update(data):
             data = data.to(device)
 
         with timeit("update"):
+            torch.compiler.cudagraph_mark_step_begin()
             loss_info = update(data)
 
         # evaluation

diff --git a/sota-implementations/iql/iql_online.py b/sota-implementations/iql/iql_online.py
@@ -176,6 +176,7 @@ def update(sampled_tensordict):
                         # sample from replay buffer
                         sampled_tensordict = replay_buffer.sample().to(device)
                     with timeit("update"):
+                        torch.compiler.cudagraph_mark_step_begin()
                         loss_info = update(sampled_tensordict)
                     # update priority
                     if prb: