The only way to speed up ops on torch is to use vectorized ops.
As @cakeeatingpolarbear mentioned, you could use torch.bmm as well as others to do so.
Perhaps you’d also like to check out tensor comprehensions 20 if you have very specialized ops not natively available on pytorch.