TensorCoding: Always use meta_ when parsing the RawMessage,

even in the case of a non-CPU resident value, because
meta_ may later be used to test for a dead Tensor.
Change: 133500231

1 files changed, 10 insertions, 3 deletions

diff --git a/tensorflow/core/distributed_runtime/tensor_coding.cc b/tensorflow/core/distributed_runtime/tensor_coding.cc
index 49507dc823..b26970b606 100644
--- a/tensorflow/core/distributed_runtime/tensor_coding.cc
+++ b/tensorflow/core/distributed_runtime/tensor_coding.cc
@@ -79,11 +79,18 @@ Status TensorResponse::ParseFrom(Source* source) {
     input.SetTotalBytesLimit(INT_MAX, INT_MAX);  // Unlimited
 
     // Pre-parse into local storage, then delegate to device.
-    RecvTensorResponse proto;
-    if (!proto.ParseFromCodedStream(&input) || !input.ConsumedEntireMessage()) {
+    if (!meta_.ParseFromCodedStream(&input) || !input.ConsumedEntireMessage()) {
       return errors::InvalidArgument("Cannot parse tensor from response");
     }
-    return device_->MakeTensorFromProto(proto.tensor(), alloc_attrs_, &tensor_);
+    Status s =
+        device_->MakeTensorFromProto(meta_.tensor(), alloc_attrs_, &tensor_);
+    // Reduce memory usage for big tensors.
+    {
+      TensorProto empty;
+      meta_.mutable_tensor()->Swap(&empty);
+    }
+    meta_.clear_tensor();
+    return s;
   }
   if (already_used_) {
     ClearTensor();