That’s what I mean by “I’m sure there is a better way to achieve this”. 
Thanks for the code. It’s a lot clearer than my approach.
I’ve called contiguous on the result, since I transposed it with t(). This might lead to errors for some operations that require contiguous tensors:
a = torch.randn(2, 4)
print(a.is_contiguous())
print(a.t().is_contiguous())