def predict(self, x):
    with torch.no_grad():
        out_raw = self.forward(x)
        out = torch.exp(out_raw)
        max_ind = out.argmax().item()
        return (self.classes[max_ind], out[:, max_ind].item())
def predict(self, x):
    with torch.no_grad():
        out_raw = self.forward(x)
        out = torch.exp(out_raw)
        max_ind = out.argmax().item()
        return (self.classes[max_ind], out[:, max_ind].item())