Merge pull request #586 from fansunqi/dim

fix dim=3 error
2024-09-05 20:18:02 +08:00 · 2024-09-05 20:18:02 +08:00 · 593519203b
parent 4b907d493c 2c141fa996
commit 593519203b
1 changed files with 1 additions and 1 deletions
--- a/python/jittor/nn.py
+++ b/python/jittor/nn.py
@ -1965,7 +1965,7 @@ def _interpolate(img, x, y, ids, mode):

 # TODO: tf_mode to another function
 def resize(img, size, mode="nearest", align_corners=False, tf_mode=False):
-    if img.dim() != 3:
+    if img.dim() != 4:
        raise ValueError("Input shape must be `(N, C, H, W)`!")
    n, c, h, w = img.shape
    H, W = size