fix fill=None workaround

pytorch · Nov 3, 2022 · 1622cd6 · 1622cd6
1 parent 3f841ca
commit 1622cd6
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/torchvision/prototype/transforms/functional/_geometry.py b/torchvision/prototype/transforms/functional/_geometry.py
@@ -691,9 +691,9 @@ def pad_image_tensor(
         )
 
     if fill is None:
-        # This is a JIT workaround
-        return _pad_with_scalar_fill(image, torch_padding, fill=0, padding_mode=padding_mode)
-    elif isinstance(fill, (int, float)):
+        fill = 0
+
+    if isinstance(fill, (int, float)):
         return _pad_with_scalar_fill(image, torch_padding, fill=fill, padding_mode=padding_mode)
     elif len(fill) == 1:
         return _pad_with_scalar_fill(image, torch_padding, fill=fill[0], padding_mode=padding_mode)