numpy
torch
pynvml
