Hi,
out from self.layer4 is 4D tensor in shape [batch, channel, h, w]. But you are refering to 3D tensor using your permute.



 krishna511:

 out = out.permute(0, 2, 1)


I am not sure about your idea, but something like out.permute(0, 2, 3, 1) will work.
PS. It would be much easier to debug if you print whole stack trace of error.
Also, documentation is clear about majority of modules’ behavior. Might need to check it.
Bests