Support Half/BFloat16 in upsample_bilinear2d (#7910)

Partial fix for #7748.
pytorch · Jan 28, 2025 · cad2d0f · cad2d0f
1 parent 8b5d1ef
commit cad2d0f
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 4 deletions.
diff --git a/kernels/portable/cpu/op_upsample_bilinear2d.cpp b/kernels/portable/cpu/op_upsample_bilinear2d.cpp
@@ -123,7 +123,7 @@ Tensor& upsample_bilinear2d_vec_out(
   const auto kernel_scale_w = area_pixel_compute_scale<double>(
       in.sizes()[3], out.sizes()[3], align_corners, scale_w);
 
-  ET_SWITCH_REAL_TYPES(
+  ET_SWITCH_REALHBF16_TYPES(
       in.scalar_type(), ctx, "upsample_bilinear2d.out", CTYPE, [&]() {
         upsample_bilinear2d_kernel_impl<CTYPE>(
             in, align_corners, kernel_scale_h, kernel_scale_w, out);

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

diff --git a/kernels/test/op_upsample_bilinear2d_test.cpp b/kernels/test/op_upsample_bilinear2d_test.cpp
@@ -302,9 +302,9 @@ TEST_F(OpUpsampleBilinear2dTest, SmokeTestAlignCornersScales) {
 }
 
 TEST_F(OpUpsampleBilinear2dTest, DType) {
-#define TEST_ENTRY(ctype, dtype)                              \
-  test_upsample_bilinear2d_dtype<ctype, ScalarType::dtype>(); \
-  ET_FORALL_REAL_TYPES(TEST_ENTRY);
+#define TEST_ENTRY(ctype, dtype) \
+  test_upsample_bilinear2d_dtype<ctype, ScalarType::dtype>();
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }