def _get_best_gpus()

in quant/common/initialization.py [0:0]

• 6 lines of code
• 3 McCabe index (conditional complexity)

def _get_best_gpus(k: int) -> List[int]:
    """Return the top k device ids associated with GPUs with the best compute capability."""
    # Select top ngpus based on CUDA device capability score
    max_gpus = torch.cuda.device_count()
    capabilities = [torch.cuda.get_device_capability(i) for i in range(max_gpus)]
    ranked_device_ids = sorted(enumerate(capabilities), key=lambda t: t[1], reverse=True)
    device_ids = [d[0] for d in ranked_device_ids][:k]
    return device_ids