diff --git a/aten/src/ATen/native/Linear.cpp b/aten/src/ATen/native/Linear.cpp index 74b5ddc6229..4ad28f3957a 100644 --- a/aten/src/ATen/native/Linear.cpp +++ b/aten/src/ATen/native/Linear.cpp @@ -70,6 +70,14 @@ static inline Tensor _flatten_nd_linear(const Tensor& input, const Tensor& weigh Tensor linear(const Tensor& input, const Tensor& weight, const c10::optional& bias_opt) { + // _matmul_impl checks this again later, but _flatten_nd_linear does not work on scalars inputs, + // so let's try to catch this here already + const auto input_dim = input.dim(); + const auto weight_dim = weight.dim(); + TORCH_CHECK(input_dim != 0 && weight_dim != 0, + "both arguments to linear need to be at least 1D, but they are ", + input_dim, "D and ", weight_dim, "D"); + // See [Note: hacky wrapper removal for optional tensor] auto bias = bias_opt.has_value() ? c10::MaybeOwned::borrowed(*bias_opt) @@ -82,7 +90,6 @@ Tensor linear(const Tensor& input, const Tensor& weight, const c10::optionaldefined()) { // Fused op is marginally faster. return at::addmm(*bias, input, weight.t()); diff --git a/test/test_nn.py b/test/test_nn.py index 23376f55741..e0c1f438fee 100644 --- a/test/test_nn.py +++ b/test/test_nn.py @@ -6526,6 +6526,14 @@ tensor(..., device='meta', size=(1,), requires_grad=True)""") expected = m(inp.view(6, 5)).view(2, 3, 8) self.assertEqual(expected, m(inp)) + def test_linear_raise_on_scalar_input(self): + # This used to cause an int underflow issue when reshaping the input + # see https://github.com/pytorch/pytorch/issues/119161 + m = nn.Linear(1, 1) + inp = torch.ones(1).squeeze() + with self.assertRaisesRegex(RuntimeError, ".*both arguments.*1D.*"): + m(inp) + @parametrize_test('device', ['cpu'] + (['cuda'] if TEST_CUDA else [])) @parametrize_test('bias', [ subtest(False, name='nobias'), subtest(True, name='bias')])