Add support for getting TensorProto argument (#18364)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/18364
att
Reviewed By: bddppq
Differential Revision: D14584784
fbshipit-source-id: 03f9207d5cf4f7f4b812428a931edbcdcb21ca8d
diff --git a/caffe2/utils/proto_utils.cc b/caffe2/utils/proto_utils.cc
index 40cc1b8..c1d6458 100644
--- a/caffe2/utils/proto_utils.cc
+++ b/caffe2/utils/proto_utils.cc
@@ -304,7 +304,14 @@
return static_cast<InputType>(static_cast<TargetType>(value)) == value;
}
}
+bool operator==(const TensorProto& l, const TensorProto& r) {
+ return l.SerializeAsString() == r.SerializeAsString();
+}
+std::ostream& operator<<(std::ostream& output, const TensorProto& n) {
+ output << n.SerializeAsString();
+ return output;
+}
bool operator==(const NetDef& l, const NetDef& r) {
return l.SerializeAsString() == r.SerializeAsString();
}
@@ -404,6 +411,7 @@
INSTANTIATE_GET_REPEATED_ARGUMENT(size_t, ints, true)
INSTANTIATE_GET_REPEATED_ARGUMENT(string, strings, false)
INSTANTIATE_GET_REPEATED_ARGUMENT(NetDef, nets, false)
+INSTANTIATE_GET_REPEATED_ARGUMENT(TensorProto, tensors, false)
#undef INSTANTIATE_GET_REPEATED_ARGUMENT
#define CAFFE2_MAKE_SINGULAR_ARGUMENT(T, fieldname) \
