fixup

intel · Jan 9, 2025 · 72dfa2a · 72dfa2a
1 parent 028d544
commit 72dfa2a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/benchmarks/triton_kernels_benchmark/gemm_postop_addmatrix_benchmark.py b/benchmarks/triton_kernels_benchmark/gemm_postop_addmatrix_benchmark.py
@@ -272,7 +272,7 @@ def matmul(a, b, d, c):
         args={},
     ))
 def benchmark(B, M, N, K, dtype, provider):
-    res_dtype = torch.float32 if dtype is torch.bfloat16 else torch.int32
+    res_dtype = torch.float32 if dtype.is_floating_point else torch.int32
     if dtype.is_floating_point:
         rand = lambda shape, dtype: torch.rand(shape, device='xpu', dtype=dtype)
     else: