| 49 | } |
| 50 | |
| 51 | Tensor* OpsTestBase::GetOutput(int output_index) { |
| 52 | CHECK_LT(output_index, context_->num_outputs()); |
| 53 | Tensor* output = context_->mutable_output(output_index); |
| 54 | #ifdef GOOGLE_CUDA |
| 55 | if (device_type_ == DEVICE_GPU) { |
| 56 | managed_outputs_.resize(context_->num_outputs()); |
| 57 | // Copy the output tensor to managed memory if we haven't done so. |
| 58 | if (!managed_outputs_[output_index]) { |
| 59 | Tensor* managed_output = |
| 60 | new Tensor(allocator(), output->dtype(), output->shape()); |
| 61 | auto src = output->tensor_data(); |
| 62 | auto dst = managed_output->tensor_data(); |
| 63 | context_->eigen_gpu_device().memcpyDeviceToHost( |
| 64 | const_cast<char*>(dst.data()), src.data(), src.size()); |
| 65 | context_->eigen_gpu_device().synchronize(); |
| 66 | managed_outputs_[output_index] = managed_output; |
| 67 | } |
| 68 | output = managed_outputs_[output_index]; |
| 69 | } |
| 70 | #endif |
| 71 | return output; |
| 72 | } |
| 73 | |
| 74 | } // namespace tensorflow |
nothing calls this directly
no test coverage detected