diff options
Diffstat (limited to 'tensorflow/core/common_runtime/gpu/gpu_device.cc')
-rw-r--r-- | tensorflow/core/common_runtime/gpu/gpu_device.cc | 8 |
1 files changed, 4 insertions, 4 deletions
diff --git a/tensorflow/core/common_runtime/gpu/gpu_device.cc b/tensorflow/core/common_runtime/gpu/gpu_device.cc index d3b6099d5b..63956afce2 100644 --- a/tensorflow/core/common_runtime/gpu/gpu_device.cc +++ b/tensorflow/core/common_runtime/gpu/gpu_device.cc @@ -114,14 +114,14 @@ class EigenCudaStreamDevice : public ::Eigen::StreamInterface { << num_bytes << ". See error logs for more detailed info."; } } - if (LogMemory::IsEnabled()) { + if (LogMemory::IsEnabled() && ret != nullptr) { LogMemory::RecordRawAllocation(operation_, step_id_, num_bytes, ret, allocator_); } return ret; } void deallocate(void* buffer) const override { - if (LogMemory::IsEnabled()) { + if (LogMemory::IsEnabled() && buffer != nullptr) { LogMemory::RecordRawDeallocation(operation_, step_id_, buffer, allocator_, true); } @@ -588,7 +588,7 @@ Status BaseGPUDeviceFactory::CreateDevices(const SessionOptions& options, for (int i = 0; i < n; i++) { BaseGPUDevice* gpu_device; TF_RETURN_IF_ERROR(CreateGPUDevice(options, - strings::StrCat(name_prefix, "/gpu:", i), + strings::StrCat(name_prefix, "/device:GPU:", i), valid_gpu_ids[i], &gpu_device)); TF_RETURN_IF_ERROR(gpu_device->Init(options)); devices->push_back(gpu_device); @@ -1049,7 +1049,7 @@ Status BaseGPUDeviceFactory::GetValidDeviceIds( size_t new_id = ids->size(); ids->push_back(visible_gpu_id); - LOG(INFO) << "Creating TensorFlow device (/gpu:" << new_id << ") -> " + LOG(INFO) << "Creating TensorFlow device (/device:GPU:" << new_id << ") -> " << "(" << GetShortDeviceDescription(visible_gpu_id, desc) << ")"; } |