Update app.py

app.py

@@ -62,17 +62,17 @@ model.load_state_dict(torch.load(model_weights_path))
 # 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.)
-# )
-
-# cnn.fit(XCnn_train, y_train)
-cnn = model
+# 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.)
+)
+
+cnn.fit(XCnn_train, y_train)
+