Update app.py

app.py

@@ -20,15 +20,27 @@ class RegressionModel2(nn.Module):
         out = self.fc2(out)
         return out
 
-# Load the trained model
-model2 = RegressionModel2(3, 32, 1)
-model2.load_state_dict(torch.load('model.pt'))
-model2.eval()
+# Load the saved model state dictionary
+model = RegressionModel2(input_dim2, hidden_dim2, output_dim2)
+model.load_state_dict(torch.load('model.pt'))
+model.eval()  # Set the model to evaluation mode
 
-def predict_astigmatism(age, aca_magnitude, aca_axis):
-    input_data = torch.tensor([[age, aca_magnitude, aca_axis]], dtype=torch.float32)
-    output = model2(input_data)
-    return output.item()
+# Define a function to make predictions
+def predict_astigmatism(age, axis, aca):
+  """
+  This function takes three arguments (age, axis, aca) as input,
+  converts them to a tensor, makes a prediction using the loaded model,
+  and returns the predicted value.
+  """
+  # Prepare the input data
+  data = torch.tensor([[age, axis, aca]], dtype=torch.float32)
+
+  # Make prediction
+  with torch.no_grad():
+    prediction = model(data)
+
+  # Return the predicted value
+  return prediction.item()
 
 def main():
     st.set_page_config(page_title='Astigmatism Prediction', page_icon=':eyeglasses:', layout='wide')