xb,yb = train_ds[i*bs : i*bs+bs]