def device(self, idx):
        """Get a specific GPU device

        Args:
            idx: index of device

        Returns:
            NvidiaDevice: single GPU device
        """
        num_dev = self.num_devices()
        assert idx < num_dev, "Cannot obtain device {}: NVML only found {} devices.".format(idx, num_dev)

        class GpuDevice(Structure):
            pass

        c_nvmlDevice_t = POINTER(GpuDevice)

        c_index = c_uint(idx)
        device = c_nvmlDevice_t()
        _check_return(_NVML.get_function(
            "nvmlDeviceGetHandleByIndex_v2")(c_index, byref(device)))
        return NvidiaDevice(device)
Get a specific GPU device

        Args:
            idx: index of device

        Returns:
            NvidiaDevice: single GPU device
        
tensorpack/utils/nvml.py