Support Half/BFloat16 in split_with_sizes.

Partial fix for #7748. ghstack-source-id: 9f183dddcd87edb2493af0f97d7ad4e40d9be434 ghstack-comment-id: 2599398274 Pull Request resolved: #7758
pytorch · Jan 18, 2025 · 779eee1 · 779eee1
1 parent 57a09f4
commit 779eee1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/kernels/portable/cpu/op_split_with_sizes_copy.cpp b/kernels/portable/cpu/op_split_with_sizes_copy.cpp
@@ -71,8 +71,8 @@ void split_with_sizes_copy_out(
   ScalarType in_type = in.scalar_type();
   ScalarType out_type = out[0].scalar_type();
 
-  ET_SWITCH_REAL_TYPES_AND(Bool, in_type, ctx, __func__, CTYPE_IN, [&]() {
-    ET_SWITCH_REAL_TYPES_AND(Bool, out_type, ctx, __func__, CTYPE_OUT, [&]() {
+  ET_SWITCH_REALHBF16_TYPES(in_type, ctx, __func__, CTYPE_IN, [&]() {
+    ET_SWITCH_REALHBF16_TYPES(out_type, ctx, __func__, CTYPE_OUT, [&]() {
       const CTYPE_IN* in_data = in.const_data_ptr<CTYPE_IN>();
 
       // Iterate through list of out tensors