for xb,yb in train_dl:
    pred = model(xb)