Running on GPU

[Gitler101]

I got this running and enjoy it, but I would like to run it on my GPU. Is this possible? I played around a bit with numba which is supposed to make it possible but I am not sure of which changes I would have to make in order to get this to work.

[CameronSima]

I believe it would depend on what hardware you're using. I'm on a MacBook Pro M1 and I changed this line:

#set device
device = torch.device('mps')
#device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

I think cuda refers to NVIDIA gpus