Define following multiplication of two tensors in pytorch lightning
I would like to to multiply following two tensors x (of shape (BS, N, C)) and y (of shape (BS,1,C)) in the following way:
BS = x.shape[0]
N = x.shape[1]
out = torch.zeros(size=x.shape)
for i in range(BS):
    for j in range(N):
        out[i, j, :] = torch.mul(x[i, j, :], y[i, 0, :])
return out


Implementing it this way yields an error &quot;RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking arugment for argument weight in method wrapper_native_layer_norm)&quot;

When fixing it by setting
out = torch.zeros(size=x.shape).to('cuda)


then training takes forever, because my for loops aren't executed in parallel.
So my question is how to implement the two for loops above in the pytorch-lightning way, so that I can define function x = multiply_as_above(x, y) and use it in the feedword(self) method of my neural network.
Btw the operation defined above looks to me like a convolution with kernel size 1. Maybe I can use that?

Is there anything wrong with x*y? As you can see in the code below, it yields exactly the same output as your function:
import torch

torch.manual_seed(2021)

BS = 2
N = 3
C = 4

x = torch.rand(BS, N, C)
y = torch.rand(BS, 1, C)

# your function
def f(x, y):
    BS = x.shape[0]
    N = x.shape[1]
    out = torch.zeros(size=x.shape)
    for i in range(BS):
        for j in range(N):
            out[i, j, :] = torch.mul(x[i, j, :], y[i, 0, :])
    return out

out1 = f(x, y)
out2 = x*y

# comparing the outputs, we can see that they are identical
torch.all(out1 == out2)
# &gt; tensor(True)

