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):

  
      
      TensorFlow
  
  
    

TensorFlow Core | Machine Learning for Beginners and Experts

Start your machine learning project with the open source ML library supported by a community of researchers, practitioners, students and developers.


  
  
    
    
  
  


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.