Currently using gpustat command, which uses nvidia-ml-py. Use the python library directly. https://pypi.org/project/nvidia-ml-py/