def gradient_descent_momentum(grads, velocity, var_list, lr, mom):
    """Gradient descent with momentum update.

    Args:
        grads: List of gradients of the trainable variables.
        velocity: List of velocity of the trainable variables.
        var_list: List of variables to be optimized.
        lr: Learning rate.
        mom: Momentum.

    Returns:
        var_list_new: List of new variables to be assigned.
        velocity_new: List of new velocity to be assigned.
    """
    velocity_new = [
        mom * v + grad
        for grad, v, var in list(zip(grads, velocity, var_list))
    ]
    var_list_new = [
        var - lr * v for v, var in list(zip(velocity_new, var_list))
    ]
    return var_list_new, velocity_new