update input0 type check for quantized FullyConnected

apache · Mar 5, 2019 · 28e1242 · 28e1242
1 parent 2f16c49
commit 28e1242
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/src/operator/quantization/quantized_fully_connected.cc b/src/operator/quantization/quantized_fully_connected.cc
@@ -79,6 +79,13 @@ bool QuantizedFullyConnectedType(const nnvm::NodeAttrs& attrs,
   CHECK_EQ(in_type->size(), num_inputs * 3);
   CHECK_EQ(out_type->size(), 3U);
 
+#if MXNET_USE_MKLDNN == 1
+  // TODO(ciyong): currently, only uint8 fully_connected is upported,
+  // int8 fully_connected will be supported after mkldnn v0.18
+  TYPE_ASSIGN_CHECK(*in_type, 0, mshadow::kUint8);
+#else
+  TYPE_ASSIGN_CHECK(*in_type, 0, mshadow::kInt8);
+#endif
   for (size_t i = 1; i < num_inputs; ++i) {
     TYPE_ASSIGN_CHECK(*in_type, i, mshadow::kInt8);
   }