Update app.py

app.py

@@ -19,11 +19,11 @@ class RegressionModel2(nn.Module):
         return out
 
 # Load the saved model state dictionaries
-model_j0 = RegressionModel2(3, 32, 1)
+model_j0 = RegressionModel2(1, 32, 1)
 model_j0.load_state_dict(torch.load('j0_model.pt'))
 model_j0.eval()
 
-model_j45 = RegressionModel2(3, 32, 1)
+model_j45 = RegressionModel2(1, 32, 1)
 model_j45.load_state_dict(torch.load('j45_model.pt'))
 model_j45.eval()
 
@@ -36,11 +36,17 @@ def calculate_initial_j0_j45(magnitude, axis_deg):
 
 def predict_new_j0_j45(age, aca_magnitude, aca_axis_deg):
     """Predict new J0 and J45 using the loaded models."""
-    aca_axis_rad = math.radians(aca_axis_deg)
-    input_data = torch.tensor([[age, aca_magnitude, aca_axis_rad]], dtype=torch.float32)
+    # Calculate initial J0 and J45
+    initial_j0, initial_j45 = calculate_initial_j0_j45(aca_magnitude, aca_axis_deg)
+    
+    # Prepare inputs for each model
+    input_data_j0 = torch.tensor([[initial_j0]], dtype=torch.float32)
+    input_data_j45 = torch.tensor([[initial_j45]], dtype=torch.float32)
+    
     with torch.no_grad():
-        new_j0 = model_j0(input_data).item()
-        new_j45 = model_j45(input_data).item()
+        new_j0 = model_j0(input_data_j0).item()
+        new_j45 = model_j45(input_data_j45).item()
+    
     return new_j0, new_j45
 
 def calculate_magnitude_and_axis(j0, j45):