summaryrefslogtreecommitdiff
path: root/caffe2/utils
diff options
context:
space:
mode:
authorYinghai Lu <yinghai@fb.com>2019-04-02 20:52:58 -0700
committerFacebook Github Bot <facebook-github-bot@users.noreply.github.com>2019-04-02 20:58:28 -0700
commit80404cb2f516ff16cefbbf9270b2003a6213b29d (patch)
tree99d1e9a534f18d515d6cca70d233ee1517ec7fde /caffe2/utils
parent31849bc524697c266552e1bd091ba9085f98494b (diff)
downloadpytorch-80404cb2f516ff16cefbbf9270b2003a6213b29d.tar.gz
pytorch-80404cb2f516ff16cefbbf9270b2003a6213b29d.tar.bz2
pytorch-80404cb2f516ff16cefbbf9270b2003a6213b29d.zip
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
Diffstat (limited to 'caffe2/utils')
-rw-r--r--caffe2/utils/proto_utils.cc8
1 files changed, 8 insertions, 0 deletions
diff --git a/caffe2/utils/proto_utils.cc b/caffe2/utils/proto_utils.cc
index 40cc1b8f31..c1d6458f92 100644
--- a/caffe2/utils/proto_utils.cc
+++ b/caffe2/utils/proto_utils.cc
@@ -304,7 +304,14 @@ bool SupportsLosslessConversion(const InputType& value) {
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(uint16_t, ints, true)
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) \