[bushishasshasha]

Description

Both Transformer implementations use output.squeeze() on a regression output of shape (batch_size, 1). When a prediction batch contains one sample, this returns a scalar, and TransformerModel.predict() fails when concatenating the resulting NumPy arrays.

Use squeeze(-1) in the DatasetH and TSDatasetH implementations to preserve the batch dimension. Multi-sample output shapes and model parameters are unchanged.

Motivation and Context

Prediction should retain every input sample, including a singleton dataset and a final batch with one remaining sample. This was reproduced on current main without market data. There is no associated issue.

How Has This Been Tested?

• Full pipeline test (pytest tests/test_all_pipeline.py from the repository root): not run. It requires external market data and broader pipeline dependencies; the official dataset is currently marked unavailable in the README. The offline tests below exercise the real prediction wrappers with Qlib datasets as a focused substitute.
• Regression tests: python -m pytest tests/model/test_transformer_batch.py -q — 12 passed. Before the fix, the same tests produced 6 failed and 6 passed.

Coverage includes both Transformer implementations, batch sizes 1/2/5, finite backward gradients, DatasetH/TSDatasetH prediction for 1/3/4 samples with batch size 2, exact output-index preservation, and agreement across inference batch partitions.

Formatting and lint:

python -m black --check -l 120 qlib/contrib/model/pytorch_transformer.py qlib/contrib/model/pytorch_transformer_ts.py tests/model/test_transformer_batch.py
python -m flake8 --ignore=E501,F541,E266,E402,W503,E731,E203 qlib/contrib/model/pytorch_transformer.py qlib/contrib/model/pytorch_transformer_ts.py tests/model/test_transformer_batch.py
git diff --check

All three checks passed locally. Tests ran on Windows with Python 3.12.6, PyTorch 2.14.1+cpu, NumPy 2.5.4, and pandas 3.1.0. Python modules were imported from this source checkout; the unchanged Cython extensions were supplied by the official pyqlib 0.9.7 Windows wheel. Full training, backtesting, and the upstream CI matrix have not been run. Existing PyTorch/pandas warnings are recorded in the test logs.

Screenshots of Test Results (if appropriate)

1. Pipeline test: not run; see limitation and offline substitute above.
2. Regression result: 12 passed, 19 warnings.

Types of changes

• Fix bugs
• Add new feature
• Update documentation

AI assistance was used to prepare the patch and regression tests.

[bushishasshasha]

@microsoft-github-policy-service agree

[Arhan Canli (arhancanli)]

The squeeze(-1) change is right for a (N, 1) decoder output, and the singleton plus remainder-batch tests cover the real failure (a final batch of one sample).

The same pattern is still on main in files this PR doesn't touch (read with gh api, not run):

• qlib/contrib/model/pytorch_localformer.py:322 and pytorch_localformer_ts.py:302: return output.squeeze() after the same self.decoder_layer(output.transpose(1, 0)[:, -1, :]).
• pytorch_gru.py:339 and pytorch_lstm.py:306: self.fc_out(out[:, -1, :]).squeeze(), also (N, 1) in, scalar out when N == 1.

So the "last batch has one sample" failure should reproduce for those models too. Would it make sense to widen this PR to them (same one-line change, and the parametrized transformer_module fixture could take their module names), or to open a follow-up so this one stays focused?