import torch._functorch.config

pytorch-labs · Dec 9, 2024 · 5217014 · 5217014
1 parent 20b138d
commit 5217014
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tritonbench/utils/triton_op.py b/tritonbench/utils/triton_op.py
@@ -621,6 +621,7 @@ def __init__(
             self.mode = Mode.BWD
         if self.mode in [Mode.FWD_BWD, Mode.BWD]:
             # TODO: remove this once we have a better way to handle backward benchmarking
+            import torch._functorch.config
             torch._functorch.config.donated_buffer = False
         self.device = tb_args.device
         self.required_metrics = (