Conv1d for two matrix column by column

I have two matrices; I need to do convolution between each column of matrix A and the correspondent column of matrix B and store the results in matrix y. This can be done using a for loop:

for i in range(5):
y[:,i]= F.conv1d(A[:, i], B[:, i])

I am wondering if there is any more efficient way to implement this process? For example, avoid using a for loop.

I don’t know if I understand your use case correctly since you did not post an executable code snippet, but assuming you want to use a “column” of B as the conv kernel, you could avoid the for loop by using a grouped convolution with a permutation in the kernel dimensions as seen here:

A = torch.randn(1, 10, 7)
B = torch.randn(1, 10, 5)
y = torch.zeros(1, 10, 3)

for i in range(A.size(1)):
    y[:, i]= F.conv1d(A[:, i:i+1], B[:, i:i+1])
out = F.conv1d(A, B.permute(1, 0, 2), groups=10)
print((out - y).abs().max())
# tensor(3.5763e-07)

Thanks, that is really helpful.