You are correct in your assumption about the missing batch dimension.
Even a single sample should contain a batch dimension with a size of 1.
Additionally to this, since you’re dealing with grayscale images (single channel), the channel dimension is also missing.
The input to nn.Conv2d
should have the shape [batch_dimension, channel, height, width]
.
In your case you could call unsqueeze
twice to add the missing dimensions of add these dimensions with a None
index:
single_loaded_img = single_loaded_img[None, None]
Also, Variables
are deprecated. If you are using a new version (> 0.3.1
), you can just remove the Variable
wrapper.