yes, you are right, after torch.cuda.synchronize() time print is normal.
the most time cost:
time of torch.matmul: 0.02192 s
one more question, does torch.matmul is very slower than conv2d ops.
my test code:
a = torch.randn(1, 153000, 1, 288)
b = torch.randn(153000, 288, 3)
#a = a.half()
#b = b.half()
if True:
device = torch.device(“cuda”)
else:
device = torch.device(“cpu”)
a = a.to(device)
b = b.to(device)
torch.cuda.synchronize()
start = time.process_time()
torch.matmul(a, b)
torch.cuda.synchronize()
end = time.process_time()
print(“Elapsed time of {}: {} s”.format(‘torch.matmul’, np.round(end - start, decimals=5)))