Update README.md

README.md

@@ -127,11 +127,11 @@ with torch.no_grad():
     images = images.to(device)
     
     for c in range(images.shape[1]):
-        # Extract single channel: (N, C, H, W) -> (N, 1, H, W)
-        single_channel = images[:, c, :, :].unsqueeze(1)
+        # Extract single channel: (N, C, H, W) -> (N, H, W)
+        single_channel = images[:, c, :, :]
         
-        # Apply transforms
-        single_channel = transform(single_channel.squeeze(1)).unsqueeze(1)
+        # Apply transforms, add dimension 1 ((N, 1, H, W))
+        single_channel = transform(single_channel).unsqueeze(1)
         
         # Extract features
         output = model.forward_features(single_channel)