Thanks for the information.
In that case you might want to use .permute, as .reshape might interleave the data:
x = torch.tensor([[0., 0.],
                  [1., 1.],
                  [2., 2.],
                  [3., 3.]])

print(x.reshape(x.size(1), x.size(0)))
> tensor([[0., 0., 1., 1.],
          [2., 2., 3., 3.]])
print(x.permute(1, 0))
> tensor([[0., 1., 2., 3.],
          [0., 1., 2., 3.]])

I don’t see any obvious errors. Could you add print statements inside the forward method, which will print the shape as well as the device of the input, output and each intermediate tensor?