add ConvRelu schema (#18693)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/18693

As title

Reviewed By: protonu

Differential Revision: D14662880

fbshipit-source-id: 3664faa660a04e1f528a413d2a1700b872c3c684

diff --git a/caffe2/quantization/server/conv_dnnlowp_op.cc b/caffe2/quantization/server/conv_dnnlowp_op.cc
index bb2aeee..d585af5 100644 (file)
--- a/caffe2/quantization/server/conv_dnnlowp_op.cc
+++ b/caffe2/quantization/server/conv_dnnlowp_op.cc
@@ -1534,9 +1534,6 @@ template class ConvDNNLowPOp<uint8_t, true>;
 template class ConvDNNLowPOp<uint16_t, false>;
 template class ConvDNNLowPOp<uint16_t, true>;
 
-OPERATOR_SCHEMA(ConvRelu).NumInputs(2, 3).NumOutputs(1).TensorInferenceFunction(
-    ConvPoolOpBase<CPUContext>::TensorInferenceForConv);
-
 REGISTER_CPU_OPERATOR_WITH_ENGINE(Conv, DNNLOWP, ConvDNNLowPOp<uint8_t, false>);
 REGISTER_CPU_OPERATOR_WITH_ENGINE(
     ConvRelu,

diff --git a/caffe2/quantization/server/conv_relu_op.cc b/caffe2/quantization/server/conv_relu_op.cc
index 6668389..f3511a8 100644 (file)
--- a/caffe2/quantization/server/conv_relu_op.cc
+++ b/caffe2/quantization/server/conv_relu_op.cc
@@ -64,6 +64,13 @@ bool ConvReluOp<T, Context>::RunOnDeviceWithOrderNHWC() {
   return true;
 }
 
+OPERATOR_SCHEMA(ConvRelu)
+    .NumInputs(2, 3)
+    .NumOutputs(1)
+    .TensorInferenceFunction(ConvPoolOpBase<CPUContext>::TensorInferenceForConv)
+    .CostInferenceFunction(OpSchema::CostInferenceFunctionType(
+        ConvPoolOpBase<CPUContext>::CostInferenceForConv));
+
 REGISTER_CPU_OPERATOR(ConvRelu, ConvReluOp<float, CPUContext>);
 
 } // namespace caffe2