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

Raises:

  • e – If CUDA runs out of memory even after lowering the batch size to 1

  • ValueError – If no output was generated

Returns:

Wrapped function

Return type:

BatchEncodingWrapper