Shahil Mawjee s-mawjee
s-mawjee
/ check_gpu.py
Last active
May 20, 2024 17:35
Get Nvidia GPU information via python code, instead of watching nvidia-smi in the terminal. Useful when training ML models, can be added to the training loop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import nvidia_smi | |
| _GPU = False | |
| _NUMBER_OF_GPU = 0 | |
| def _check_gpu(): | |
| global _GPU | |
| global _NUMBER_OF_GPU | |
| nvidia_smi.nvmlInit() | |
| _NUMBER_OF_GPU = nvidia_smi.nvmlDeviceGetCount() |