train on gpus?

[indazole]

Hello,
Thank you for posting this cool work!
Wondering if you are able to train using a GPU?
Without getting into details, I have only been able to train with CPUs.
Cheers-

[eloyfelix]

this is quite a shallow neural network, so I'm unsure using GPU can improve the training a lot. I'll take a look to make cpu/gpu selectable but I'm not very confident you'll see an speed increase here.