Update handler.py

handler.py

@@ -1,43 +1,31 @@
 import torch
 import torch.nn as nn
-import torch.optim as optim
-import os
 from huggingface_hub import PyTorchModelHubMixin
 
+# Set device
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  
+print('Device:', device)
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # Set device
-print('device >>> ', device)
+# Define model class
+class MyModel(nn.Module, PyTorchModelHubMixin):
+    def __init__(self):
+        super().__init__()  # Initialize nn.Module
+        # Model layers will be loaded via from_pretrained()
 
+    def forward(self, x):
+        return self.model(x)  # Assume this model has a defined forward pass
 
-
-
-
-class EndpointHandler():
-
-
-    # Define model class that loads a pretrained model
-    class MyModel(nn.Module, PyTorchModelHubMixin):
-       def __init__(self, model_name="damiano216/pay-boo-2"):
-           super().__init__()
-           # Load pretrained model from Hugging Face
-           self.model = MyModel.from_pretrained(model_name)
- 
-       def forward(self, x):
-           return self.model(x)
-
-
-    net = MyModel()
-
+# EndpointHandler class
+class EndpointHandler:
     def __init__(self, path=""):
-        self.model = MyModel()  # Load the pretrained model
-        self.model.to(device)  # Move to GPU if available
-        self.model.eval()  # Set to evaluation mode
-
+        self.model = MyModel.from_pretrained("damiano216/pay-boo-2")  # Load from Hugging Face
+        self.model.to(device)  # Move model to GPU if available
+        self.model.eval()  # Set model to evaluation mode
 
     def __call__(self, data):
-        new_data_tensor = torch.tensor(data['chargeData'], dtype=torch.float32).to(device)  # Convert input to tensor
+        new_data_tensor = torch.tensor(data['chargeData'], dtype=torch.float32).to(device)  # Ensure tensor is on device
 
-        # 3. Make predictions
+        # Make predictions
         with torch.no_grad():
             predictions = self.model(new_data_tensor)