import mindspore as ms
from mindspore import Tensor, ops
import numpy as np
ms.set_context(mode=ms.GRAPH_MODE)
input_x = Tensor(np.arange(6480000).reshape((32, 450, 450)), ms.float32)
mat2 = Tensor(np.arange(6480000).reshape((32, 450, 450)), ms.float32)
output = ops.bmm(input_x, mat2)
print(output)
print(output.shape)

batmatmul = ops.BatchMatMul()
output = batmatmul(input_x, mat2)
print(output)
print(output.shape)