[Relay] add a dimension check to reject invalid input (#14925)

* add a dimension check

* Update test_forward.py

* Update pytorch.py

* Update pytorch.py

python/tvm/relay/frontend/pytorch.py

@@ -4642,6 +4642,12 @@ def get_relay_ty(ishape, itype, pt_type):
             ):
                 msg = "Shapes of input list and information in the graph do not match"
                 raise RuntimeError(msg)
+            if len(ishape) > 1 and any(dim <= 0 for dim in ishape[1:]):
+                msg = (
+                    "Expected input's non-batch dimensions to have positive length, "
+                    f"but input has a shape of {pt_type.sizes()}"
+                )
+                raise RuntimeError(msg)
             pt_dtype = pt_type.scalarType()
             if not pt_dtype and itype:
                 pt_dtype = itype

tests/python/frontend/pytorch/test_forward.py

@@ -789,6 +789,11 @@ def test_forward_celu():
     input_data = torch.tensor([-1.0, 2.0], dtype=torch.float32)
     verify_model(torch.nn.CELU().eval(), input_data=input_data)
 
+    input_shape = [2, 0, 1]
+    input_data = torch.rand(input_shape).float()
+    with pytest.raises(RuntimeError):
+        verify_model(torch.nn.CELU().eval(), input_data=input_data)
+
 
 @tvm.testing.uses_gpu
 def test_forward_gelu():