fixed rocm batch_matmul strategy for mixed i8i8i32

masahi · Apr 13, 2022 · 762c7e8 · 762c7e8
1 parent ce53e8d
commit 762c7e8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/tvm/relay/op/strategy/rocm.py b/python/tvm/relay/op/strategy/rocm.py
@@ -212,7 +212,7 @@ def batch_matmul_strategy_rocm(attrs, inputs, out_type, target):
     """Batch matmul strategy for ROCM"""
     strategy = _op.OpStrategy()
     strategy.add_implementation(
-        wrap_compute_batch_matmul(topi.cuda.batch_matmul),
+        wrap_compute_batch_matmul(topi.cuda.batch_matmul, need_out_dtype=True),
         wrap_topi_schedule(topi.cuda.schedule_batch_matmul),
         name="batch_matmul.cuda",
         plevel=10,