Regression bug?

I tried the regression option with 4000+ rows but it returns multiple Y (predicted) columns and the output fails here:

/content/tabfm/tabfm/src/classifier_and_regressor.py in _predict_internal(self, X)
3285
3286 output = self._batch_forward(Xs_all, ys_all, cat_masks_all, ds=ds_all)
-> 3287 predictions = output.squeeze(-1)
3288 return predictions
3289

Because it expects a single Y output.

I ran a debug version to look at the returned data structure and this is how I discovered the multiple Y output issue.

Reproduce the regression case with 4000+ rows, then inspect classifier_and_regressor.py at _predict_internal and the _batch_forward output it receives. Confirm the expected prediction shape for regression and that the failing case completes without the multiple-Y output error.

Written by the indexing model from the issue text.