For that, you can use cupy[0], PyTorch[1] or Tensorflow[2]. They all mimic the numpy's API with the possibility to use your GPU.

[0] https://cupy.dev/ [1] https://pytorch.org/ [2] https://www.tensorflow.org/