batch_encoding_wrapper
- lightning_ir.base.model.batch_encoding_wrapper(func: BatchEncodingWrapper) BatchEncodingWrapper[source]
Decorator to enable sub-batching for models that support it. Lowers the batch size of the input batch encoding if the model runs out of memory.
- Parameters:
func (BatchEncodingWrapper) – Function to wrap that takes a batch encoding.
- Returns:
Wrapped function that handles sub-batching.
- Return type:
- Raises:
RuntimeError – If CUDA runs out of memory and the batch size cannot be lowered further.
ValueError – If no output was generated.