Update app.py

app.py

@@ -21,7 +21,10 @@ class Net(nn.Module):
 
 # Load the trained model
 model = Net()
-model.load_state_dict(torch.load('mnist_model.pth', map_location=torch.device('cpu')))
+try:
+    model.load_state_dict(torch.load('mnist_model.pth', map_location=torch.device('cpu')))
+except Exception as e:
+    print(f"Error loading model: {e}")
 model.eval()
 
 # Define the transform to preprocess the input image
@@ -51,4 +54,3 @@ iface = gr.Interface(
 # Launch the Gradio interface
 if __name__ == "__main__":
     iface.launch()
-