disable gemm activation for non-float data types

microsoft · Feb 22, 2024 · e8772c7 · e8772c7
1 parent 76a2a48
commit e8772c7
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/onnxruntime/core/optimizer/gemm_activation_fusion.cc b/onnxruntime/core/optimizer/gemm_activation_fusion.cc
@@ -65,6 +65,12 @@ Status GemmActivationFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l
       continue;
     }
 
+    NodeArg* node_output = node.MutableOutputDefs()[0];
+    auto data_type = node_output->TypeAsProto()->tensor_type().elem_type();
+    if (data_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT) {
+      continue;
+    }
+
     Node& gemm_node = node;
     Node& act_node = *graph.GetNode(next_node.Index());  // get mutable reference