diff options
Diffstat (limited to 'tensorflow/contrib/verbs/rdma.h')
| -rw-r--r-- | tensorflow/contrib/verbs/rdma.h | 7 |
1 files changed, 7 insertions, 0 deletions
diff --git a/tensorflow/contrib/verbs/rdma.h b/tensorflow/contrib/verbs/rdma.h index 10cbbe58d9..16ef58bc62 100644 --- a/tensorflow/contrib/verbs/rdma.h +++ b/tensorflow/contrib/verbs/rdma.h @@ -28,6 +28,7 @@ limitations under the License. #include <vector> #include "tensorflow/core/distributed_runtime/worker_env.h" +#include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/platform/env.h" @@ -225,6 +226,12 @@ class RdmaTensorBuffer : public RdmaBuffer { explicit RdmaTensorBuffer(RdmaChannel* channel, string name); virtual ~RdmaTensorBuffer() override {} void SendNextItem() override; + void PostCopyOperations(bool can_memcpy, size_t buffer_size, + size_t tensor_bytes, const string& key, + const Tensor& in, int64 step_id, bool is_dead, + const string& key_with_step_id, const Tensor* copy, + const TensorProto* proto, + const StringPiece* copy_buf); }; struct RdmaMessage { |