import torch
import torch.nn.functional as F

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
    
    @classmethod
    def load(cls, path, device=None):
        device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        model = cls()
        model.load_state_dict(torch.load(path, map_location=device))
        model.to(device)
        model.eval()
        return model

class Inference:
    def __init__(self, model_path):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model = Model.load(model_path, device=self.device)

    def run(self, input_data):
        if not isinstance(input_data, torch.Tensor):
            input_data = torch.tensor(input_data, dtype=torch.float32)
        
        input_data = input_data.to(self.device)
        
        with torch.no_grad():
            output = self.model(input_data)
            return torch.argmax(output, dim=-1)

if __name__ == "__main__":
    predictor = Inference("path/to/model.pth")
    prediction = predictor.run([0.1, 0.2, 0.3])