Fix bug on autocontrast when min==max

Author: datumbox

test/test_functional_tensor.py

@@ -820,6 +820,18 @@ def test_autocontrast(device, dtype, channels):
     )
 
 
+@pytest.mark.parametrize("device", cpu_and_gpu())
+@pytest.mark.parametrize("dtype", (None, torch.float32, torch.float64))
+@pytest.mark.parametrize("channels", [1, 3])
+def test_autocontrast_equal_minmax(device, dtype, channels):
+    a = _create_data_batch(32, 32, num_samples=1, channels=channels, device=device)
+    a = a / 2.0 + 0.3
+    assert (F.autocontrast(a)[0] == F.autocontrast(a[0])).all()
+
+    a[0, 0] = 0.7
+    assert (F.autocontrast(a)[0] == F.autocontrast(a[0])).all()
+
+
 @pytest.mark.parametrize("device", cpu_and_gpu())
 @pytest.mark.parametrize("channels", [1, 3])
 def test_equalize(device, channels):

torchvision/transforms/functional_tensor.py

@@ -937,10 +937,10 @@ def autocontrast(img: Tensor) -> Tensor:
 
     minimum = img.amin(dim=(-2, -1), keepdim=True).to(dtype)
     maximum = img.amax(dim=(-2, -1), keepdim=True).to(dtype)
-    eq_idxs = torch.where(minimum == maximum)[0]
-    minimum[eq_idxs] = 0
-    maximum[eq_idxs] = bound
     scale = bound / (maximum - minimum)
+    eq_idxs = torch.isfinite(scale).logical_not()
+    minimum[eq_idxs] = 0
+    scale[eq_idxs] = 1
 
     return ((img - minimum) * scale).clamp(0, bound).to(img.dtype)