Update app.py

app.py

@@ -52,7 +52,13 @@ torch.manual_seed(0)
 
 
 # Create an instance of your model
-model = Cnn()
+model = NeuralNetClassifier(
+    Cnn,
+    max_epochs=10,
+    lr=0.002,
+    optimizer=torch.optim.Adam,
+    device=device,
+)
 
 # Specify the path to the saved model weights
 model_weights_path = 'model_weights.pth'
@@ -62,6 +68,14 @@ model.load_state_dict(torch.load(model_weights_path,map_location=torch.device('c
 
 # Set the model to evaluation mode for inference
 model.eval()
+# Create a NeuralNetClassifier using the loaded model
+cnn = NeuralNetClassifier(
+    module=model,
+    max_epochs=0,  # Set max_epochs to 0 to avoid additional training
+    lr=0.002,  # You can set this to the learning rate used during training
+    optimizer=torch.optim.Adam,  # You can set the optimizer used during training
+    device='cpu'  # You can specify the device ('cpu' for CPU, 'cuda' for GPU, etc.)
+)
 
 stroke_width = st.sidebar.slider("Stroke width: ", 1, 35, 32)
 stroke_color = st.sidebar.color_picker("Stroke color hex: ")
@@ -97,6 +111,6 @@ if canvas_result.image_data is not None:
     st.image(image1)
 
     image1.resize(1,1,28,28)
-    st.title(np.argmax(model.predict(image1)))
+    st.title(np.argmax(cnn.predict(image1)))
 if canvas_result.json_data is not None:
     st.dataframe(pd.json_normalize(canvas_result.json_data["objects"]))