add softmax

app.py

@@ -4,6 +4,7 @@ import torch
 from torch import nn
 import numpy as np
 from PIL import Image
+import torch.nn.functional as F
 
 # Creating Neural Model class used for training and validation
 class NeuralNetwork(nn.Module):
@@ -21,6 +22,7 @@ class NeuralNetwork(nn.Module):
     def forward(self, x):
         x = self.flatten(x)
         logits = self.linear_relu_stack(x)
+        logits = F.log_softmax(logits, dim=1)
         return logits
 
 # Intializing and loading the saved model