Support Half/BFloat16 in unbind_copy (#7908)

Partial fix for #7748.
pytorch · Jan 28, 2025 · fe86377 · fe86377
1 parent a85a42b
commit fe86377
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/kernels/portable/cpu/op_unbind_copy.cpp b/kernels/portable/cpu/op_unbind_copy.cpp
@@ -54,10 +54,10 @@ void unbind_copy_int_out(
   ScalarType in_type = input.scalar_type();
   ScalarType out_type = out[0].scalar_type();
 
-  ET_SWITCH_REAL_TYPES_AND(
-      Bool, in_type, ctx, "unbind_copy.int_out", CTYPE_IN, [&]() {
-        ET_SWITCH_REAL_TYPES_AND(
-            Bool, out_type, ctx, "unbind_copy.int_out", CTYPE_OUT, [&]() {
+  ET_SWITCH_REALHBF16_TYPES(
+      in_type, ctx, "unbind_copy.int_out", CTYPE_IN, [&]() {
+        ET_SWITCH_REALHBF16_TYPES(
+            out_type, ctx, "unbind_copy.int_out", CTYPE_OUT, [&]() {
               const CTYPE_IN* const input_data =
                   input.const_data_ptr<CTYPE_IN>();
               for (size_t i = 0, e = out.size(); i < e; ++i) {

diff --git a/kernels/test/op_unbind_copy_test.cpp b/kernels/test/op_unbind_copy_test.cpp
@@ -208,19 +208,19 @@ class OpUnbindCopyIntOutTest : public OperatorTest {
  */
 TEST_F(OpUnbindCopyIntOutTest, Unbind1x2x3OnDim0AllRealDtypes) {
 #define TEST_ENTRY(ctype, dtype) test_unbind_dim0<ScalarType::dtype>();
-  ET_FORALL_REAL_TYPES(TEST_ENTRY);
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }
 
 TEST_F(OpUnbindCopyIntOutTest, Unbind1x2x3OnDim1AllRealDTypes) {
 #define TEST_ENTRY(ctype, dtype) test_unbind_dim1<ScalarType::dtype>();
-  ET_FORALL_REAL_TYPES(TEST_ENTRY);
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }
 
 TEST_F(OpUnbindCopyIntOutTest, Unbind1x2x3OnDim2AllRealDTypes) {
 #define TEST_ENTRY(ctype, dtype) test_unbind_dim2<ScalarType::dtype>();
-  ET_FORALL_REAL_TYPES(TEST_ENTRY);
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }