From 89e9b1cf8e869799b14e5673fa0631982059c671 Mon Sep 17 00:00:00 2001
From: Jongsoo Park <jongsoo@fb.com>
Date: Mon, 1 Apr 2019 13:02:02 -0700
Subject: [PATCH] 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
---
 caffe2/quantization/server/conv_dnnlowp_op.cc | 3 ---
 caffe2/quantization/server/conv_relu_op.cc    | 7 +++++++
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/caffe2/quantization/server/conv_dnnlowp_op.cc b/caffe2/quantization/server/conv_dnnlowp_op.cc
index bb2aeee..d585af5 100644
--- 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
--- 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
-- 
2.7.4