Update app.py

app.py

@@ -1,10 +1,26 @@
 import os
+import tempfile
+import numpy as np
+import torch
+import torch.nn as nn
 import onnx
 from onnx import external_data_helper
 
-def export_onnx_model(trained_model: LitClassifier, mu: np.ndarray, sd: np.ndarray, n_features: int) -> str:
-    wrapper = OnnxWrapper(trained_model.net.cpu().eval(), mu=mu, sd=sd).eval()
+class OnnxWrapper(nn.Module):
+    def __init__(self, net: nn.Module, mu: np.ndarray, sd: np.ndarray):
+        super().__init__()
+        self.net = net
+        self.register_buffer("mu", torch.tensor(mu, dtype=torch.float32))
+        self.register_buffer("sd", torch.tensor(sd, dtype=torch.float32))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = (x - self.mu) / self.sd
+        logits = self.net(x).squeeze(-1)
+        return torch.sigmoid(logits)
 
+def export_onnx_model(trained_model, mu: np.ndarray, sd: np.ndarray, n_features: int) -> str:
+    # trained_model is your LightningModule; we export its .net
+    wrapper = OnnxWrapper(trained_model.net.cpu().eval(), mu=mu, sd=sd).eval()
     dummy = torch.zeros(1, n_features, dtype=torch.float32)
 
     tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".onnx", prefix="model_")
@@ -21,13 +37,12 @@ def export_onnx_model(trained_model: LitClassifier, mu: np.ndarray, sd: np.ndarr
         do_constant_folding=True,
     )
 
-    # If exporter wrote external data, merge it back into a single .onnx
+    # Merge external data back into one .onnx if needed
     data_path = onnx_path + ".data"
     if os.path.exists(data_path):
         m = onnx.load_model(onnx_path, load_external_data=True)
         external_data_helper.convert_model_from_external_data(m)
         onnx.save_model(m, onnx_path)
-        # remove external blob now that it's embedded
         try:
             os.remove(data_path)
         except OSError: