| 205 | } |
| 206 | |
| 207 | bool GetDeviceName(TFE_Context* ctx, string* device_name, |
| 208 | const char* device_type) { |
| 209 | std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status( |
| 210 | TF_NewStatus(), TF_DeleteStatus); |
| 211 | TF_DeviceList* devices = TFE_ContextListDevices(ctx, status.get()); |
| 212 | CHECK_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); |
| 213 | |
| 214 | const int num_devices = TF_DeviceListCount(devices); |
| 215 | for (int i = 0; i < num_devices; ++i) { |
| 216 | const string dev_type(TF_DeviceListType(devices, i, status.get())); |
| 217 | CHECK_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get()); |
| 218 | const string dev_name(TF_DeviceListName(devices, i, status.get())); |
| 219 | CHECK_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get()); |
| 220 | if (dev_type == device_type) { |
| 221 | *device_name = dev_name; |
| 222 | LOG(INFO) << "Found " << device_type << " device " << *device_name; |
| 223 | TF_DeleteDeviceList(devices); |
| 224 | return true; |
| 225 | } |
| 226 | } |
| 227 | TF_DeleteDeviceList(devices); |
| 228 | return false; |
| 229 | } |
no test coverage detected