There’s a lot of weird things here, but to address the error you are getting:
This particular error should be gone if you change this line:



 Leo_Frez:

x = torch.randn(480, 640).view(-1, 3, 480, 640) # batch, channel, height, width
 



To either this:
x = torch.randn(3, 480, 640).view(-1, 3, 480, 640) # batch, channel, height, width

Or this:
x = torch.randn(1, 3, 480, 640) # batch, channel, height, width