enable_chunking

enable_chunking(module_or_flag=None, chunk_size=None)

Enable chunked FFT conv globally, as decorator, or as context manager.

Can be used as: 1. A function to set global state: enable_chunking(True) 2. A class decorator: @enable_chunking on an nn.Module 3. A context manager: with enable_chunking():

Parameters:

• module_or_flag – Either a bool, module class, or None.

• chunk_size (int | None) – Optional chunk size to set.

Returns:

Decorated class, context manager, or None depending on usage.