diff --git a/models.py b/models.py
index c90eeba7..8af1368a 100644
--- a/models.py
+++ b/models.py
@@ -218,7 +218,7 @@ def _basic_init(module):
     def unpatchify(self, x):
         """
         x: (N, T, patch_size**2 * C)
-        imgs: (N, H, W, C)
+        imgs: (N, C, H, W)
         """
         c = self.out_channels
         p = self.x_embedder.patch_size[0]