pynvml
matplotlib
torch