class matmul_transpose (torch.autograd.Function):
使用NPU自定义算子替换原生写法,以提高性能。
from torch_npu.contrib.function import matmul_transpose output = matmul_transpose(tensor1, tensor2)
>>> tensor1 = torch.randn(68, 5, 75, 16).npu() >>> tensor1.requires_grad_(True) >>> tensor2 = torch.randn(68, 5, 75, 16).npu() >>> tensor2.requires_grad_(True) >>> output = matmul_transpose(tensor1, tensor2) >>> output.sum().backward()