It looks like you are passing the input as “channel last”.
Currently you have to pass image tensors as [batch_size, channels, height, width].
You could use x.permute(0, 3, 1, 2) to permute the dimensions.