Albert_GZ:

image = np.array(image).astype(np.float32)
    image /= 255.
    image = torch.from_numpy(image).to(cuda)



Hi,
The error is about the channel logic. numpy is channel last, pytorch is channel first, the error says you are using numpy logic for pytorch tensors. Instead of converting your PIL images to numpy then reading those arrays as tensor using from_numpy, you can directly convert the PIL image to proper tensor using transforms.ToTensor. Other workaround would be permuting channels of the image after converting it to tensor using from_numpy by using torch.permute.
The idea is you should get [1, 3, 720, 840], instead of [720, 840, 3, 1].
Bests