Update app.py

app.py

@@ -1,5 +1,6 @@
 # Import necessary libraries
 import numpy as np
+import os
 import gradio as gr
 import torch
 import torch.nn as nn
@@ -53,6 +54,18 @@ iface = gr.Interface(
     description="Upload an image of a handwritten digit, and the model will predict the digit."
 )
 
+# Check if the file exists
+if not os.path.isfile('mnist_model.pth'):
+    raise FileNotFoundError("The model file 'mnist_model.pth' was not found.")
+else:
+    print("Model file found, proceeding with loading.")
+
+# Load the model state dict
+model.load_state_dict(torch.load('mnist_model.pth'))
+
+model.load_state_dict(torch.load('mnist_model.pth', weights_only=True))
+
+
 # Launch the Gradio interface
 if __name__ == '__main__':
     iface.launch()