use the correct datatype format (#8144)

seravee · ezyang · commit ffde23d45eee · 2018-06-05T22:01:59.000-04:00
diff --git a/caffe2/operators/shape_op.cc b/caffe2/operators/shape_op.cc
@@ -22,7 +22,7 @@ OPERATOR_SCHEMA(Shape)
       } else {
         out[0].add_dims(axes.size());
       }
-      out[0].set_data_type(TensorProto::INT32);
+      out[0].set_data_type(TensorProto::INT64);
       return out;
     })
     .SetDoc(R"DOC(
diff --git a/caffe2/python/operator_test/shape_inference_test.py b/caffe2/python/operator_test/shape_inference_test.py
@@ -494,6 +494,12 @@ def testInt8Conversion(self):
         workspace.FeedBlob('x', np.random.rand(100, 150).astype(np.float32))
         self.InferTensorRunAndCompare(model)
 
+    def testShapeOp(self):
+        model = model_helper.ModelHelper(name="shape_op_test")
+        model.Shape('x', 'y')
+        workspace.FeedBlob('x', np.random.rand(100, 150).astype(np.float32))
+        self.InferTensorRunAndCompare(model)
+
     def InferTensorRunAndCompare(self, model, expected_uninferred_blobs=None):
         '''
         Runs shape inference, and then the model to check

Original file line number	Diff line number	Diff line change
`@@ -22,7 +22,7 @@ OPERATOR_SCHEMA(Shape)`
`22`	`22`	`} else {`
`23`	`23`	`out[0].add_dims(axes.size());`
`24`	`24`	`}`
`25`		`- out[0].set_data_type(TensorProto::INT32);`
	`25`	`+ out[0].set_data_type(TensorProto::INT64);`
`26`	`26`	`return out;`
`27`	`27`	`})`
`28`	`28`	`.SetDoc(R"DOC(`