From 21fbc78e69b3e79f883161601f1ccc62be42a06e Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Thu, 26 Jul 2018 23:07:53 -0700 Subject: Set the correct context when calling cudnnCreate. When running with multiple devices, using the wrong context will lead to a check-fail when trying to set a stream that has been created with a different context. This resolves a check-fail on resnet50 with 8 GPUs. PiperOrigin-RevId: 206274741 --- tensorflow/stream_executor/cuda/cuda_dnn.cc | 1 + 1 file changed, 1 insertion(+) (limited to 'tensorflow/stream_executor') diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.cc b/tensorflow/stream_executor/cuda/cuda_dnn.cc index 766a0dafb5..1c3940e92c 100644 --- a/tensorflow/stream_executor/cuda/cuda_dnn.cc +++ b/tensorflow/stream_executor/cuda/cuda_dnn.cc @@ -322,6 +322,7 @@ port::Status GetLoadedCudnnVersion(CudnnVersion* version) { CudnnSupport::CudnnSupport(CUDAExecutor* parent) : parent_(parent) {} port::Status CudnnSupport::Init() { + ScopedActivateExecutorContext context(parent_); cudnnHandle_t cudnn_handle = nullptr; auto status = cudnnCreate(&cudnn_handle); if (status == CUDNN_STATUS_SUCCESS) { -- cgit v1.2.3