Update test/test_prototype_transforms.py

Author: pmeier

test/test_prototype_transforms.py

@@ -197,9 +197,7 @@ def input_expected_image_tensor(self, p, dtype=torch.float32):
         input = torch.tensor([[[0, 1], [0, 1]], [[1, 0], [1, 0]]], dtype=dtype)
         expected = torch.tensor([[[1, 0], [1, 0]], [[0, 1], [0, 1]]], dtype=dtype)
 
-        if p == 1.0:
-            return input, expected
-        return input, input
+        return input, expected if p == 1 else input
 
     def test_simple_tensor(self, p):
         input, expected = self.input_expected_image_tensor(p)