import torch.nn.functional as F

loss_func = F.cross_entropy

def model(xb):
    return xb @ weights + bias