Fix gemv

pytorch-labs · Nov 18, 2024 · 597facf · 597facf
1 parent 59bb8fe
commit 597facf
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 2 deletions.
diff --git a/test/test_gpu/main.py b/test/test_gpu/main.py
@@ -37,8 +37,6 @@
 
 TEST_OPERATORS = set(list_operators_by_collection(op_collection="default")) - SKIP_OPS
 
-print(f"Testing operators: {TEST_OPERATORS}")
-
 
 def check_ci_output(op):
     from tritonbench.utils.triton_op import REGISTERED_BENCHMARKS

diff --git a/tritonbench/operators/addmm/hstu.py b/tritonbench/operators/addmm/hstu.py
@@ -9,6 +9,7 @@
 with add_path(str(SUBMODULE_PATH.joinpath("generative-recommenders"))):
     from generative_recommenders.ops.triton.triton_addmm import _addmm_fwd
 
+
 class _AddMmFunction(torch.autograd.Function):
     @staticmethod
     # pyre-ignore[14]

diff --git a/tritonbench/utils/triton_op.py b/tritonbench/utils/triton_op.py
@@ -521,11 +521,13 @@ def __call__(cls, *args, **kwargs):
         obj.__post__init__()
         return obj
 
+
 def _translate_mode(tb_args):
     def _has_and_true(attr):
         if hasattr(tb_args, attr) and getattr(tb_args, attr):
             return True
         return False
+
     if _has_and_true("fwd"):
         tb_args.mode = "fwd"
     if _has_and_true("bwd"):
@@ -535,6 +537,7 @@ def _has_and_true(attr):
     if _has_and_true("fwd_no_grad"):
         tb_args.mode = "fwd_no_grad"
 
+
 class BenchmarkOperator(metaclass=PostInitProcessor):
     mode: Mode = Mode.FWD
     test: str = "eval"