Drop the leading batch dimension for single-sequence demo outputs.
(key, value)
| 266 | |
| 267 | |
| 268 | def _squeeze_single_batch(key, value): |
| 269 | """Drop the leading batch dimension for single-sequence demo outputs.""" |
| 270 | batched_ndim = _BATCHED_NDIMS.get(key) |
| 271 | if batched_ndim is None or not hasattr(value, "ndim"): |
| 272 | return value |
| 273 | if value.ndim == batched_ndim and value.shape[0] == 1: |
| 274 | return value[0] |
| 275 | return value |
| 276 | |
| 277 | |
| 278 | def postprocess(predictions, images): |
no outgoing calls
no test coverage detected