Simple, complete GPU NN example

I’m testing GPU acceleration with AI frameworks. I am not familiar with PyTorch. I’m looking for a “hello world” example that trains and tests a neural network, and uses the GPU when available. To see what I have in mind, check out this MNIST example with Tensorflow (the left-hand panel):

It’s small, it’s complete and in one piece, it does something interesting, and will use a GPU when available.

Strangely, I was not able to find something similar with PyTorch. I’ve seen examples where changes were suggested to make it run on the GPU (but assumed familiarity with PyTorch), I’ve seen fragments of code (e.g. just the training part), I’ve seen very large pieces of code - but never something short, complete, using a neural network, and GPU-enabled.

Could someone please point me at the right example? Thank you.

This tuto is maybe what you are looking for

1 Like

That was perfect, thank you!

1 Like