diff --git a/tensorflow/core/framework/tensor.cc b/tensorflow/core/framework/tensor.cc index 683dac05d2f..e5eb512a642 100644 --- a/tensorflow/core/framework/tensor.cc +++ b/tensorflow/core/framework/tensor.cc @@ -650,6 +650,12 @@ Tensor::Tensor(DataType type, const TensorShape& shape, TensorBuffer* buf) RefIfNonNull(buf); } +Tensor::Tensor(DataType type, const TensorShape& shape, + core::RefCountPtr<TensorBuffer> buf) + : shape_(shape), buf_(buf.release()) { + set_dtype(type); +} + bool Tensor::IsInitialized() const { return (buf_ != nullptr && buf_->data() != nullptr) || shape_.num_elements() == 0; diff --git a/tensorflow/core/framework/tensor.h b/tensorflow/core/framework/tensor.h index 062e5413c55..33a240d85dc 100644 --- a/tensorflow/core/framework/tensor.h +++ b/tensorflow/core/framework/tensor.h @@ -157,6 +157,12 @@ class Tensor { /// Acquires a ref on buf that belongs to this Tensor. Tensor(DataType type, const TensorShape& shape, TensorBuffer* buf); + /// \brief Creates a tensor with the input datatype, shape and buf. + /// + /// Takes an ownership of the bufffer from the reference counted pointer. + Tensor(DataType type, const TensorShape& shape, + core::RefCountPtr<TensorBuffer> buf); + /// \brief Creates an empty Tensor of the given data type. /// /// Like Tensor(), returns a 1-dimensional, 0-element Tensor with