diff options
| author |  A. Unique TensorFlower <gardener@tensorflow.org> | 2018-07-26 23:07:53 -0700 |
|---|---|---|
| committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-07-26 23:10:45 -0700 |
| commit | 21fbc78e69b3e79f883161601f1ccc62be42a06e (patch) | |
| tree | 9438003d582541359062700cbb9e7f0e95a7a894 /tensorflow/stream_executor | |
| parent | 19defc68c99483049e14ec18718cb037f1ca050b (diff) | |
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
Diffstat (limited to 'tensorflow/stream_executor')
| -rw-r--r-- | tensorflow/stream_executor/cuda/cuda_dnn.cc | 1 |
1 files changed, 1 insertions, 0 deletions
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) { |