TODO verify the dataloader. ensure the data is being batched correctly edit the model to operate on numpy arrays. predictions should be am array of shape (bs, 10) verify the parameters are updating properly.