def METHOD_NAME(self, set_to_none=None):
    if self.gradient_state.sync_gradients:
        accept_arg = "set_to_none" in inspect.signature(self.optimizer.METHOD_NAME).parameters
        if accept_arg:
            if set_to_none is None:
                set_to_none = False
            self.optimizer.METHOD_NAME(set_to_none=set_to_none)
        else:
            if set_to_none is not None:
                raise ValueError("`set_to_none` for Optimizer.zero_grad` is not supported by this optimizer.")
            self.optimizer.METHOD_NAME()