Fix unit test

pytorch-labs · Dec 2, 2024 · f0e2a04 · f0e2a04
1 parent 3f75a34
commit f0e2a04
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tritonbench/operators/layer_norm/operator.py b/tritonbench/operators/layer_norm/operator.py
@@ -40,6 +40,7 @@ def torch_compile_layer_norm(self, *args):
             import torch._functorch.config
 
             torch._functorch.config.donated_buffer = False
+        import torch
 
         @torch.compile
         def inner(*args):