File tree Expand file tree Collapse file tree 2 files changed +7
-2
lines changed Expand file tree Collapse file tree 2 files changed +7
-2
lines changed Original file line number Diff line number Diff line change @@ -136,7 +136,12 @@ Status MatMul<T>::Compute(OpKernelContext* ctx) const {
136136 // be filled out with zeros.
137137 EigenMatrixMapRowMajor<T> dest (y->MutableData <T>(),
138138 narrow<Eigen::Index>(helper.M ()), narrow<Eigen::Index>(helper.N ()));
139- dest.setZero ();
139+ if constexpr (std::is_same<T, MLFloat16>::value) {
140+ dest.setConstant (MLFloat16 (0 .0f ));
141+ } else {
142+ dest.setZero ();
143+ }
144+
140145 return Status::OK ();
141146 }
142147
Original file line number Diff line number Diff line change @@ -311,7 +311,7 @@ TEST(MathOpTest, MatMul_Float16) {
311311 run_test (true );
312312 run_test (false );
313313}
314- // #endif
314+ #endif
315315
316316#if defined(USE_CUDA) || defined(USE_ROCM) || defined(USE_DNNL)
317317TEST (MathOpTest, MatMul_bfloat16) {
You can’t perform that action at this time.
0 commit comments