Fix compilation of contrib.mpi and contrip.mpi_collectives when build…

…ing with CUDA support
tensorflow · Mar 14, 2019 · 03d0080 · 03d0080
1 parent bacd851
commit 03d0080
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/tensorflow/contrib/mpi/mpi_rendezvous_mgr.cc b/tensorflow/contrib/mpi/mpi_rendezvous_mgr.cc
@@ -29,6 +29,7 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/gpu/gpu_util.h"
 #include "tensorflow/core/distributed_runtime/session_mgr.h"
 #include "tensorflow/core/distributed_runtime/tensor_coding.h"
+#include "tensorflow/core/framework/allocator.h"
 
 namespace tensorflow {
 
@@ -122,7 +123,7 @@ void MPIRemoteRendezvous::RecvFromRemoteAsync(
         } else {
           TensorResponse tr;
           tr.InitAlloc(dst_device, recv_args.alloc_attrs);
-          tr.InitPartial(mpi_response.response());
+          tr.InitPartial(mpi_response.response(), AllocationAttributes());
           const size_t nBytes = tr.tensor().TotalBytes();
           void* data = const_cast<void*>(DMAHelper::base(&tr.tensor()));
           MPI_Status status;

diff --git a/tensorflow/contrib/mpi_collectives/kernels/ring.cu.cc b/tensorflow/contrib/mpi_collectives/kernels/ring.cu.cc
@@ -19,6 +19,8 @@ limitations under the License.
 
 #define EIGEN_USE_GPU
 
+#include "tensorflow/core/util/cuda_launch_config.h"
+
 #include "tensorflow/contrib/mpi_collectives/kernels/ring.h"
 
 namespace tensorflow {