Update mnist_model.pth

mnist_model.pth

@@ -1,2 +1,5 @@
-# Save the trained model state dictionary
+# train_model.py
+import torch
+
+# Assuming `model` is your trained model
 torch.save(model.state_dict(), 'mnist_model.pth')