src/transformers/models/esm/openfold_utils/chunk_utils.py/ChunkSizeTuner/__init__
class ChunkSizeTuner:    def __init__(
        self,
        # Heuristically, runtimes for most of the modules in the network
        # plateau earlier than this on all GPUs I've run the model on.
        max_chunk_size: int = 512,
    ):
        self.max_chunk_size = max_chunk_size
        self.cached_chunk_size: Optional[int] = None
        self.cached_arg_data: Optional[tuple] = None