I think it would be much simpler with the functional interface. Just call F.conv2d(input, weight) where weight is generated by some other part of the network. This should work with arbitrary batches, just be careful about providing correct dimensions of the weights.