diff --git a/utils/datasets.py b/utils/datasets.py
index abb4a3650bfc..d3714d745b88 100755
--- a/utils/datasets.py
+++ b/utils/datasets.py
@@ -574,7 +574,7 @@ def __getitem__(self, index):
             labels_out[:, 1:] = torch.from_numpy(labels)
 
         # Convert
-        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
+        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3 x img_height x img_width
         img = np.ascontiguousarray(img)
 
         return torch.from_numpy(img), labels_out, self.img_files[index], shapes