def METHOD_NAME(self, model: nn.Module, should_step: bool):
    if should_step:
        yield
    else:
        # Prevents DDP from syncing gradients during accumulation step
        with model.no_sync():
            yield