in src/sagemaker_training/environment.py [0:0]
def num_gpus(): # type: () -> int
"""Return the number of GPUs available in the current container.
Returns:
int: Number of GPUs available in the current container.
"""
try:
cmd = shlex.split("nvidia-smi --list-gpus")
output = subprocess.check_output(cmd).decode("utf-8")
return sum([1 for x in output.split("\n") if x.startswith("GPU ")])
except (OSError, subprocess.CalledProcessError):
logger.info("No GPUs detected (normal if no gpus installed)")
return 0