diff --git a/kernels/portable/cpu/op_flip.cpp b/kernels/portable/cpu/op_flip.cpp index fab4792a5a4..41e99953c93 100644 --- a/kernels/portable/cpu/op_flip.cpp +++ b/kernels/portable/cpu/op_flip.cpp @@ -66,7 +66,7 @@ Tensor& flip_out( constexpr auto name = "flip.out"; - ET_SWITCH_REALHB_TYPES(in.scalar_type(), ctx, name, CTYPE, [&] { + ET_SWITCH_REALHBBF16_TYPES(in.scalar_type(), ctx, name, CTYPE, [&] { const CTYPE* in_data = in.const_data_ptr(); CTYPE* out_data = out.mutable_data_ptr(); diff --git a/kernels/test/op_flip_test.cpp b/kernels/test/op_flip_test.cpp index 01c79a283e9..3301a243b56 100644 --- a/kernels/test/op_flip_test.cpp +++ b/kernels/test/op_flip_test.cpp @@ -33,20 +33,26 @@ class OpFlipOutTest : public ::testing::Test { // first. torch::executor::runtime_init(); } + + template + void test_1d_dtype() { + TensorFactory tf; + + Tensor input = tf.make({4, 1, 3}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); + int64_t dims_data[1] = {-1}; + IntArrayRef dims = IntArrayRef(dims_data, 1); + Tensor out = tf.zeros({4, 1, 3}); + Tensor out_expected = + tf.make({4, 1, 3}, {3, 2, 1, 6, 5, 4, 9, 8, 7, 12, 11, 10}); + op_flip_out(input, dims, out); + EXPECT_TENSOR_CLOSE(out, out_expected); + } }; TEST_F(OpFlipOutTest, SmokeTest1Dim) { - TensorFactory tfFloat; - - Tensor input = - tfFloat.make({4, 1, 3}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); - int64_t dims_data[1] = {-1}; - IntArrayRef dims = IntArrayRef(dims_data, 1); - Tensor out = tfFloat.zeros({4, 1, 3}); - Tensor out_expected = - tfFloat.make({4, 1, 3}, {3, 2, 1, 6, 5, 4, 9, 8, 7, 12, 11, 10}); - op_flip_out(input, dims, out); - EXPECT_TENSOR_CLOSE(out, out_expected); +#define TEST_ENTRY(ctype, dtype) test_1d_dtype(); + ET_FORALL_REALHBF16_TYPES(TEST_ENTRY); +#undef TEST_ENTRY } TEST_F(OpFlipOutTest, SmokeTest2Dims) {