def log_softmax(x):
    return x - x.exp().sum(-1).log().unsqueeze(-1)

def model([xb](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor")):
    return log_softmax([xb](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor") @ [weights](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor") + [bias](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor"))