Update (base update)

[ghstack-poisoned]
pytorch · Dec 13, 2024 · 836b83c · 836b83c
1 parent 8920180
commit 836b83c
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()