diff --git a/utils/segment/general.py b/utils/segment/general.py
index 9da8945..f1b2f1d 100644
--- a/utils/segment/general.py
+++ b/utils/segment/general.py
@@ -10,7 +10,7 @@ def crop_mask(masks, boxes):
     Vectorized by Chong (thanks Chong).
 
     Args:
-        - masks should be a size [h, w, n] tensor of masks
+        - masks should be a size [n, h, w] tensor of masks
         - boxes should be a size [n, 4] tensor of bbox coords in relative point form
     """