Support Half/BFloat16 in upsample_nearest2d (pytorch#7911)

Partial fix for pytorch#7748.
zonglinpeng · Jan 30, 2025 · 42659f8 · 42659f8
1 parent 457b5a8
commit 42659f8
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/kernels/portable/cpu/op_upsample_nearest2d.cpp b/kernels/portable/cpu/op_upsample_nearest2d.cpp
@@ -79,7 +79,7 @@ Tensor& upsample_nearest2d_vec_out(
   const auto kernel_scale_w = area_pixel_compute_scale<double>(
       in.sizes()[3], out.sizes()[3], false, scale_w);
 
-  ET_SWITCH_REAL_TYPES(
+  ET_SWITCH_REALHBF16_TYPES(
       in.scalar_type(), ctx, "upsample_nearest2d.out", CTYPE, [&]() {
         upsample_nearest2d_kernel_impl<CTYPE>(
             in, kernel_scale_h, kernel_scale_w, out);

diff --git a/kernels/test/CMakeLists.txt b/kernels/test/CMakeLists.txt
@@ -230,6 +230,7 @@ set(all_test_sources
     "op_unbind_copy_test.cpp"
     "op_unsqueeze_copy_test.cpp"
     "op_upsample_bilinear2d_test.cpp"
+    "op_upsample_nearest2d_test.cpp"
     "op_var_test.cpp"
     "op_view_copy_test.cpp"
     "op_where_test.cpp"

diff --git a/kernels/test/op_upsample_nearest2d_test.cpp b/kernels/test/op_upsample_nearest2d_test.cpp
@@ -52,7 +52,6 @@ class OpUpsampleNearest2dTest : public OperatorTest {
     op_upsample_nearest2d_out(
         input,
         OptionalArrayRef<int64_t>({output_size.data(), output_size.size()}),
-        true,
         {},
         out);
 
@@ -254,9 +253,9 @@ TEST_F(OpUpsampleNearest2dTest, MultiBatchAndChannel) {
 }
 
 TEST_F(OpUpsampleNearest2dTest, DType) {
-#define TEST_ENTRY(ctype, dtype)                             \
-  test_upsample_nearest2d_dtype<ctype, ScalarType::dtype>(); \
-  ET_FORALL_REAL_TYPES(TEST_ENTRY);
+#define TEST_ENTRY(ctype, dtype) \
+  test_upsample_nearest2d_dtype<ctype, ScalarType::dtype>();
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }