title: Define following multiplication of two tensors in pytorch lightning
body: <p>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:</p>
  <pre><code>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
  </code></pre>
  <ol>
  <li><p>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;</p>
  </li>
  <li><p>When fixing it by setting</p>
  <p>out = torch.zeros(size=x.shape).to('cuda)</p>
  </li>
  </ol>
  <p>then training takes forever, because my for loops aren't executed in parallel.</p>
  <p>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?</p>