def _format_param(name, optimizer, param):
    """Return correctly formatted lr/momentum for each param group."""
    if isinstance(param, (list, tuple)):
        if len(param) != len(optimizer.param_groups):
            raise ValueError("expected {} values for {}, got {}".format(
                len(optimizer.param_groups), name, len(param)))
        return param
    else:
        return [param] * len(optimizer.param_groups)