Multiclass classification using pytorch

Double post from here with potential answer.