Solved by eqy in post #2 
              
              
                I think this is a natural place to do an einsum as it is effectively a variation of an outer-product style computation: 
import torch
import time

B = 16
C = 64
H = 32
W = 32

input = torch.randn(B, 3, H, W)
mask = torch.randn(B, C, H, W)

output = torch.zeros(B*C,3,H,W)
t1 = time.time()
for b in ra…