| 23 | namespace tensorflow { |
| 24 | |
| 25 | void OpsTestBase::SetDevice(const DeviceType& device_type, |
| 26 | std::unique_ptr<Device> device) { |
| 27 | CHECK(device_) << "No device provided"; |
| 28 | |
| 29 | device_ = device.get(); |
| 30 | device_mgr_ = absl::make_unique<DeviceMgr>(std::move(device)); |
| 31 | pflr_ = absl::make_unique<ProcessFunctionLibraryRuntime>( |
| 32 | device_mgr_.get(), Env::Default(), TF_GRAPH_DEF_VERSION, flib_def_.get(), |
| 33 | OptimizerOptions()); |
| 34 | |
| 35 | device_type_ = device_type; |
| 36 | #ifdef GOOGLE_CUDA |
| 37 | if (device_type == DEVICE_GPU) { |
| 38 | managed_allocator_.reset(new GpuManagedAllocator()); |
| 39 | allocator_ = managed_allocator_.get(); |
| 40 | } else { |
| 41 | managed_allocator_.reset(); |
| 42 | allocator_ = device_->GetAllocator(AllocatorAttributes()); |
| 43 | } |
| 44 | #else |
| 45 | CHECK_NE(device_type, DEVICE_GPU) |
| 46 | << "Requesting GPU on binary compiled without GOOGLE_CUDA."; |
| 47 | allocator_ = device_->GetAllocator(AllocatorAttributes()); |
| 48 | #endif |
| 49 | } |
| 50 | |
| 51 | Tensor* OpsTestBase::GetOutput(int output_index) { |
| 52 | CHECK_LT(output_index, context_->num_outputs()); |
no test coverage detected