Hi,
You are using wrong permutation order.
See this example:
x = torch.randn(1, 128, 128, 3)
# your order
x.permute(0,2,3,1).shape # torch.Size([1, 128, 3, 128])
# correct order
x = x.permute(0, 3, 1, 2)
x.shape # torch.Size([1, 3, 128, 128])
And the error corresponds to this issue.
I am not sure still it could work fine or not because of channel mutation, because of the forward method. For instance, concatenation, squeezing, etc and all other method which use dim argument to do the operation, if exist in the forward function, may cause issues. You may need to override forward function w.r.t. channel changes.
Bests