Hey all,
I have a tensor t with shape (b,c,n,m) where b is the batch size, c is the number of channels, n is the sequence length (number of tokens) and m a number of parallel representations of the data (similar to the different heads in the transformer).
I want to perform a 1D conv over the channels and sequence length, such that each block would have its own convolution layer. pseudo-code:
t = torch.rand([b,c,n,m])
convs = [conv1d(c,c,1) for _ in range(m)]
for i in range(m)
   output[:,:,:,i] = convs[i](input[:,:,:,i])

I’m pretty sure that for loop isn’t the way to go. How can I perform this computation efficiently?
Thanks!