Update app.py

app.py

@@ -20,12 +20,12 @@ class LaserPredictions(nn.Module):
         return out
 
 # Load the saved model state dictionary
-model = LaserPredictions(6, 32, 3)
+model = LaserPredictions(3, 32, 3)
 model.load_state_dict(torch.load('laser_prescription_model.pt'))
 model.eval()  # Set the model to evaluation mode
 
-def predict(pre_op_sphere, pre_op_cylinder, pre_op_axis, three_month_sphere, three_month_cylinder, three_month_axis):
-    input_data = torch.tensor([[pre_op_sphere, pre_op_cylinder, pre_op_axis, three_month_sphere, three_month_cylinder, three_month_axis]], dtype=torch.float32)
+def predict(pre_op_sphere, pre_op_cylinder, pre_op_axis):
+    input_data = torch.tensor([[pre_op_sphere, pre_op_cylinder, pre_op_axis]], dtype=torch.float32)
     
     with torch.no_grad():
         predicted_prescription = model(input_data)
@@ -40,9 +40,6 @@ inputs = [
     gr.Number(label="Pre-Op Sphere"),
     gr.Number(label="Pre-Op Cylinder"),
     gr.Number(label="Pre-Op Axis"),
-    gr.Number(label="3-Month Sphere"),
-    gr.Number(label="3-Month Cylinder"),
-    gr.Number(label="3-Month Axis")
 ]
 output = gr.Textbox(label="Predicted Laser Prescription")