[FIX] Fix cublas batch matmul (apache#6715)

* Update batch_matmul.py Update batch_matmul.py * fix
trevor-m · Dec 4, 2020 · f8f17a3 · f8f17a3
1 parent 7607ade
commit f8f17a3
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/python/tvm/topi/cuda/batch_matmul.py b/python/tvm/topi/cuda/batch_matmul.py
@@ -138,7 +138,7 @@ def _callback(op):
     return s
 
 
-def batch_matmul_cublas(x, y):
+def batch_matmul_cublas(x, y, out_shape=None):
     """Computes batch matrix multiplication of `x` and `y` when `x` and `y` are
     data in batch.
 
@@ -150,6 +150,9 @@ def batch_matmul_cublas(x, y):
     y : tvm.te.Tensor
         3-D with shape [batch, N, K]
 
+    out_shape : None
+        The output shape
+
     Returns
     -------
     output : tvm.te.Tensor