Solved by ptrblck in post #2 
              
              
                It seems your input is missing the batch dimension. 
Generally you should pass the input as [batch_size, channels, height, width]. 
In your use case, if you load a single image it will be in the shape [c, h, w]. 
Just unsqueeze the input and pass it to the model: 
img_data = img_data.unsqueeze(0)
se…