Make _is_tensor_a_torch_image more generic

fmassa · fmassa · commit 016784b15824 · 2020-06-01T15:01:12.000+02:00
diff --git a/torchvision/transforms/functional_tensor.py b/torchvision/transforms/functional_tensor.py
@@ -4,7 +4,7 @@
 
 
 def _is_tensor_a_torch_image(input):
-    return len(input.shape) == 3
+    return input.ndim >= 2
 
 
 def vflip(img):
