yb = y_train[0:bs]
print(loss_func(preds, yb))