add input check at the beginning

heidongxianhua · heidongxianhua · commit 79e6a86b3804 · 2023-09-04T15:28:45.000+08:00
diff --git a/test/cpp/api/functional.cpp b/test/cpp/api/functional.cpp
@@ -2236,6 +2236,13 @@ TEST_F(FunctionalTest, Interpolate) {
       }
     }
   }
+  {
+    ASSERT_THROWS_WITH(
+        F::interpolate(
+            torch::randn({1}),
+            F::InterpolateFuncOptions().size(std::vector<int64_t>({1}))),
+        "Input Error: Only 3D, 4D and 5D input Tensors supported (got 1D) ");
+  }
   {
     auto input = torch::randn({3, 2, 2});
     ASSERT_THROWS_WITH(
diff --git a/torch/csrc/api/include/torch/nn/functional/upsampling.h b/torch/csrc/api/include/torch/nn/functional/upsampling.h
@@ -107,6 +107,16 @@ inline Tensor interpolate(
     }
   }
 
+  TORCH_CHECK(
+      input.dim() >= 3 && input.dim() <= 5,
+      "Input Error: Only 3D, 4D and 5D input Tensors supported "
+      "(got ",
+      input.dim(),
+      "D) for the modes: nearest | linear | bilinear | bicubic | trilinear "
+      "(got ",
+      enumtype::get_enum_name(mode),
+      ")");
+
   auto scale_factor_len = input.dim() - 2;
   std::vector<c10::optional<double>> scale_factor_list(
       scale_factor_len, c10::nullopt);
