Author: Goran Trlin
Multiclass classification demo in PyTorch is an example of using a small neural network written in PyTorch to classify points in 2D space. Based on its coordinates, each point can be assigned to one of the four predefined classes. The presented neural network is a sequence of linear layers enhanced with a ReLu layer. Once trained, the network is set to inference mode and tested on a test dataset. The example can be run either on a CUDA enabled GPU or on CPU (default option).
The full source code is located here.