bugfix

pytorch-labs · Nov 14, 2024 · 0d856f9 · 0d856f9
1 parent 12f6551
commit 0d856f9
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 2 deletions.
diff --git a/tritonbench/operators/ragged_attention/hstu.py b/tritonbench/operators/ragged_attention/hstu.py
@@ -150,7 +150,7 @@ def forward(
             grid = (1216,)
             _ragged_hstu_attn_fwd_persistent[grid](**kwargs)
         else:
-            _RaggedAttentionRelativeBiasFunction.apply(
+            out = _RaggedAttentionRelativeBiasFunction.apply(
                 self.max_seq_len,  # N
                 kwargs["alpha"],
                 q,

diff --git a/tritonbench/operators/ragged_attention/operator.py b/tritonbench/operators/ragged_attention/operator.py
@@ -82,7 +82,6 @@ def get_bwd_fn(self, fwd_fn: Callable[..., Any]) -> Callable[..., Any]:
             lambda x: isinstance(x, torch.Tensor),
             o,
         )
-        print(o)
         do = torch.rand_like(o_tensor)
         fn = lambda: o_tensor.backward(do, retain_graph=True)
         return fn