torch_module.py

import torch



def matmul_func(a, b):
    y = torch.broadcast_to(b, (a.shape[0], b.shape[0], b.shape[1]))
    y = torch.transpose(y, 1, 2).float()
    z = torch.bmm(a, y)
    return z



compiled_module = torch.jit.script(matmul_func)

print(compiled_module.graph)

a = torch.randn(4, 128, 3200)
b = torch.randn(8640, 3200).to(torch.half)

import time


for i in range(10):
    start = time.perf_counter()
    x = compiled_module(a, b)
    end = time.perf_counter()
    print(f"The time taken for the operation is:  {(end - start)*1000} ms")