Hi,
The reshape operation is differentiable in pytorch. So it should not prevent you from computing any gradient.
Also I think you want to permute dimensions, not reshape them: input = input.permute(1, 0, 2, 3).