This is an automated email from the ASF dual-hosted git repository.
echuraev pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new d776bccc1e [Relay] add a dimension check to reject invalid input
(#14925)
d776bccc1e is described below
commit d776bccc1edb90b06d06b4ea4a94157115b1d972
Author: Qingchao Shen <[email protected]>
AuthorDate: Wed May 24 13:28:17 2023 +0800
[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 | 6 ++++++
tests/python/frontend/pytorch/test_forward.py | 5 +++++
2 files changed, 11 insertions(+)
diff --git a/python/tvm/relay/frontend/pytorch.py
b/python/tvm/relay/frontend/pytorch.py
index 08bf5d517c..f3bcda7c23 100644
--- a/python/tvm/relay/frontend/pytorch.py
+++ b/python/tvm/relay/frontend/pytorch.py
@@ -4642,6 +4642,12 @@ def _get_relay_input_vars(graph, input_infos, prelude,
is_module=True, default_d
):
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
diff --git a/tests/python/frontend/pytorch/test_forward.py
b/tests/python/frontend/pytorch/test_forward.py
index b602c14df3..de1de6421b 100644
--- a/tests/python/frontend/pytorch/test_forward.py
+++ b/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():