diff options
author | Yinghai Lu <yinghai@fb.com> | 2019-04-02 20:52:58 -0700 |
---|---|---|
committer | Facebook Github Bot <facebook-github-bot@users.noreply.github.com> | 2019-04-02 20:58:28 -0700 |
commit | 80404cb2f516ff16cefbbf9270b2003a6213b29d (patch) | |
tree | 99d1e9a534f18d515d6cca70d233ee1517ec7fde /caffe2/utils | |
parent | 31849bc524697c266552e1bd091ba9085f98494b (diff) | |
download | pytorch-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.cc | 8 |
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) \ |