Update

[ghstack-poisoned]
pytorch · Dec 13, 2024 · 2a802dd · 2a802dd
2 parents 163ca84 + 836b83c
commit 2a802dd
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/sota-implementations/decision_transformer/online_dt.py b/sota-implementations/decision_transformer/online_dt.py
@@ -136,6 +136,7 @@ def update(data):
             data = offline_buffer.sample()
 
         with timeit("update"):
+            torch.compiler.cudagraph_mark_step_begin()
             loss_vals = update(data.to(model_device))
 
         scheduler.step()