Validate matching input shapes in Int8Add operator (#14520)

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

Default engine doesn't support broadcast semantics in Int8Add operator. This patch adds a check that shapes are equivalent.

Reviewed By: bertmaher

Differential Revision: D13250922

fbshipit-source-id: 8526d07723bd9a34d54dee04d121c57f8b33c481

diff --git a/caffe2/operators/quantized/int8_add_op.h b/caffe2/operators/quantized/int8_add_op.h
index 5689880..4f30195 100644 (file)
--- a/caffe2/operators/quantized/int8_add_op.h
+++ b/caffe2/operators/quantized/int8_add_op.h
@@ -32,6 +32,11 @@ class Int8AddOp final : public Operator<CPUContext> {
     const auto& B = Inputs()[1]->template Get<Int8TensorCPU>();
     auto* Y = Outputs()[0]->template GetMutable<Int8TensorCPU>();
 
+    CAFFE_ENFORCE_EQ(
+        A.t.sizes(),
+        B.t.sizes(),
+        "inputs must have the same shape (broadcast semantics is not supported)");
+
     /*
      * Record quantization parameters for A and B inputs, because if the op is
      * in-place, we may overwrite these parameters later, when we set