File size: 520 Bytes
ab687e7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 |
import torch
import warnings
def check_gpus_available(ngpus: int) -> None:
ngpus_available = torch.cuda.device_count()
if ngpus < ngpus_available:
msg = 'Not using all available GPUS.' + \
f' N GPUs available: {ngpus_available},' + \
f' N GPUs selected: {ngpus}. '
warnings.warn(msg)
elif ngpus > ngpus_available:
msg = 'Not enough GPUs to satisfy selected amount' + \
f': {ngpus}. N GPUs available: {ngpus_available}'
warnings.warn(msg)
|