Add non-regression test

pytorch · Feb 6, 2024 · 3100374 · 3100374
1 parent b8f5610
commit 3100374
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/test/test_transforms_v2.py b/test/test_transforms_v2.py
@@ -5182,6 +5182,11 @@ def test_functional_and_transform(self, make_input, fn):
         if isinstance(input, torch.Tensor):
             assert output.data_ptr() == input.data_ptr()
 
+    def test_2d_np_array(self):
+        # Non-regression test for https://github.com/pytorch/vision/issues/8255
+        input = np.random.rand(10, 10)
+        assert F.to_image(input).shape == (1, 10, 10)
+
     def test_functional_error(self):
         with pytest.raises(TypeError, match="Input can either be a pure Tensor, a numpy array, or a PIL image"):
             F.to_image(object())