model = Mnist_Logistic()
print(loss_func(model(xb), yb))