import torch

class Inference:
    def __init__(self, model_path, device=None):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        
        # We set weights_only=False because we trust our own Model.save format
        self.model = torch.load(model_path, map_location=self.device, weights_only=False)
        
        self.model.to(self.device)
        self.model.eval()

    def run(self, input_data):
        with torch.no_grad():
            if not isinstance(input_data, torch.Tensor):
                input_data = torch.tensor(input_data, dtype=torch.float32)
            
            # Ensure input has a batch dimension if it's missing
            if input_data.ndim == 1:
                input_data = input_data.unsqueeze(0)
                
            input_data = input_data.to(self.device)
            output = self.model(input_data)
            return torch.argmax(output, dim=1)