Update MNIST Classification Model..py

Added Logging for Loss vs Epoch

Author: infinitecoder1729

MNIST Classification Model..py

@@ -12,7 +12,6 @@
 
 train_batch = tch.utils.data.DataLoader(train, batch_size=30, shuffle=True)
 
-
 input = 784
 hidden = 490
 output = 10
@@ -40,13 +39,16 @@
         optimize.zero_grad()
         output = model(images)
         loss = lossfn(output, labels)
+        writer.add_scalar("Loss", loss, num)
         loss.backward()
         optimize.step()
         run += loss.item()
     else:
         print("Epoch Number : {} = Loss : {}".format(num, run/len(train_batch)))
 Elapsed=(time()-time_start)/60
 print("\nTraining Time (in minutes) : ",Elapsed)
+writer.flush()
+writer.close()
 
 correct=0
 all = 0