diff --git a/kernels/optimized/cpu/op_log_softmax.cpp b/kernels/optimized/cpu/op_log_softmax.cpp index d49e72efcc2..a9117216ac6 100644 --- a/kernels/optimized/cpu/op_log_softmax.cpp +++ b/kernels/optimized/cpu/op_log_softmax.cpp @@ -102,6 +102,7 @@ void log_softmax_kernel(const Tensor& input, int64_t dim, Tensor& out) { template bool log_softmax_wrapper(const Tensor& X, int64_t dim, Tensor& out) { if constexpr ( + std::is_same_v || std::is_same_v || std::is_same_v) { // Input dtype equals output dtype (enforced by check_log_softmax_args). @@ -112,7 +113,6 @@ bool log_softmax_wrapper(const Tensor& X, int64_t dim, Tensor& out) { } else { auto input_scalar_type = X.scalar_type(); switch (input_scalar_type) { - // TODO: support Double as well case ScalarType::Float: log_softmax_kernel(X, dim, out); return true; @@ -149,12 +149,16 @@ Tensor& opt_log_softmax_out( auto out_scalar_type = out.scalar_type(); switch (out_scalar_type) { - // TODO: support Double as well case ScalarType::Float: { bool success = log_softmax_wrapper(self, dim, out); ET_KERNEL_CHECK(context, success, InvalidArgument, out); break; } + case ScalarType::Double: { + bool success = log_softmax_wrapper(self, dim, out); + ET_KERNEL_CHECK(context, success, InvalidArgument, out); + break; + } case ScalarType::BFloat16: { bool success = log_softmax_wrapper(self, dim, out); diff --git a/kernels/optimized/test/supported_features_def.yaml b/kernels/optimized/test/supported_features_def.yaml index df51f3ec2f4..054ba32e29a 100644 --- a/kernels/optimized/test/supported_features_def.yaml +++ b/kernels/optimized/test/supported_features_def.yaml @@ -1,5 +1,2 @@ - namespace: op_gelu dtype_double: false - -- namespace: op_log_softmax - dtype_double: false