Solved by eqy in post #2 
              
              
                I think you can express this computation as either a batched matrix multiply or an einsum that implements a batched matrix multiply: 
import time
import torch

x = torch.randn(64, 3, 64, 64, dtype=torch.double)
test = torch.randn(32, 48)
def orig(x):
    out = torch.zeros(x.size(1), x.size(1), dtype…