| 31 | namespace tensorflow { |
| 32 | |
| 33 | static std::vector<std::unique_ptr<BaseGPUDevice>> GetGPUDevices() { |
| 34 | std::vector<std::unique_ptr<Device>> devices; |
| 35 | SessionOptions session_options; |
| 36 | session_options.config.mutable_gpu_options() |
| 37 | ->set_per_process_gpu_memory_fraction(0.1); |
| 38 | session_options.env = Env::Default(); |
| 39 | Status s = DeviceFactory::GetFactory(DEVICE_GPU) |
| 40 | ->AddDevices(session_options, "", &devices); |
| 41 | TF_CHECK_OK(s); |
| 42 | std::vector<std::unique_ptr<BaseGPUDevice>> gpus; |
| 43 | for (std::unique_ptr<Device>& device : devices) { |
| 44 | if (device->device_type() == "GPU") { |
| 45 | // If `device_type()` is GPU, this `Device` is guaranteed to be a |
| 46 | // `BaseGPUDevice`, which is a subclass of `Device`. |
| 47 | gpus.emplace_back(static_cast<BaseGPUDevice*>(device.release())); |
| 48 | } |
| 49 | } |
| 50 | return gpus; |
| 51 | } |
| 52 | |
| 53 | template <typename Scalar> |
| 54 | class NcclManagerTest : public ::testing::Test { |
no test coverage detected