Make this code as base:
import math

import matplotlib.pyplot as plt
import numpy as np
import torch
from hyperopt import fmin, hp, tpe

import torch
from torch.optim import Optimizer
import warnings
import math

plt.style.use("seaborn-v0_8-white")


def rosenbrock(tensor):
    # https://en.wikipedia.org/wiki/Test_functions_for_optimization
    x, y = tensor
    return (1 - x) ** 2 + 100 * (y - x ** 2) ** 2


def execute_steps(func, initial_state, optimizer_class, optimizer_config, num_iter=500):
    x = torch.Tensor(initial_state).requires_grad_(True)
    optimizer = optimizer_class([x], **optimizer_config)
    steps = []
    steps = np.zeros((2, num_iter + 1))
    steps[:, 0] = np.array(initial_state)
    for i in range(1, num_iter + 1):
        optimizer.zero_grad()
        f = func(x)
        f.backward()
        torch.nn.utils.clip_grad_norm_(x, 1.0)
        optimizer.step()
        steps[:, i] = x.detach().numpy()
    return steps


def objective_rosenbrok(params):
    lr = params["lr"]
    optimizer_class = params["optimizer_class"]
    minimum = (1.0, 1.0)
    initial_state = (-2.0, 2.0)
    optimizer_config = dict(lr=lr)
    num_iter = 400
    steps = execute_steps(rosenbrock, initial_state, optimizer_class, optimizer_config, num_iter)
    return (steps[0][-1] - minimum[0]) ** 2 + (steps[1][-1] - minimum[1]) ** 2


def plot_rosenbrok(grad_iter, optimizer_name, lr):
    x = np.linspace(-2, 2, 250)
    y = np.linspace(-1, 3, 250)
    minimum = (1.0, 1.0)

    X, Y = np.meshgrid(x, y)
    Z = rosenbrock([X, Y])

    iter_x, iter_y = grad_iter[0, :], grad_iter[1, :]

    fig = plt.figure(figsize=(8, 8))

    ax = fig.add_subplot(1, 1, 1)
    ax.contour(X, Y, Z, 90, cmap="jet")
    ax.plot(iter_x, iter_y, color="r", marker="x")

    ax.set_title("Rosenbrock func: {} with {} "
                 "iterations, lr={:.6}".format(optimizer_name, len(iter_x), lr))
    plt.plot(*minimum, "gD")
    plt.plot(iter_x[-1], iter_y[-1], "rD")
    plt.savefig("../docs/rosenbrock_{}.png".format(optimizer_name))


def execute_experiments(optimizers, objective, func, plot_func, initial_state, seed=1):
    seed = seed
    for item in optimizers:
        optimizer_class, lr_low, lr_hi = item
        space = {"optimizer_class": hp.choice("optimizer_class", [optimizer_class]),
                 "lr": hp.loguniform("lr", lr_low, lr_hi), }
        best = fmin(fn=objective, space=space, algo=tpe.suggest, max_evals=300, rstate=np.random.default_rng(seed))
        print(best["lr"], optimizer_class)

        steps = execute_steps(func, initial_state, optimizer_class, {"lr": best["lr"]}, num_iter=500, )
        plot_func(steps, optimizer_class.__name__, best["lr"])


class Lilith(Optimizer):
    def __init__(self, params, lr: float, eps: float = 1e-8, beta1_m: float = 0.9, beta2_m: float = 0.9,
                 beta_v: float = 0.999, weight_decay: float = 0., m_norm_min: float = 1e-4, ratio_min: float = 1e-4,
                 lookahead_k: int = 5, lookahead_beta: float = 0.5):
        defaults = dict(lr=lr, eps=eps, beta1_m=beta1_m, beta2_m=beta2_m, beta_v=beta_v, weight_decay=weight_decay,
                        m_norm_min=m_norm_min, ratio_min=ratio_min, lookahead_k=lookahead_k,
                        lookahead_beta=lookahead_beta)

        super(Lilith, self).__init__(params, defaults)

    @torch.no_grad
    def step(self, closure=None):
        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:
            for p in group['params']:
                if p.grad is None:
                    continue

                grad = p.grad.data

                state = self.state[p]

                if len(state) == 0:
                    state['step'] = 0
                    state['m_avg1'] = torch.zeros_like(grad)
                    state['m_avg2'] = torch.zeros_like(grad)
                    state['v_avg'] = torch.zeros_like(grad)
                    state['ema'] = p.data.clone()

                state['step'] += 1

                if sum(grad.shape) > 1:
                    trust_ratio = (p.data.norm() / grad.norm().clip(min=1e-4)).clip(min=group['ratio_min'])
                    grad.sub_(grad.mean(dim=tuple(range(1, len(grad.shape))), keepdim=True))
                    grad.mul_(trust_ratio)

                m_avg1_prev = state['m_avg1'].clone()
                state['m_avg1'].add_(state['m_avg2']).lerp_(grad, 1 - group['beta1_m'])
                state['m_avg2'].lerp_(state['m_avg1'] - m_avg1_prev, 1 - group['beta2_m'])

                u = state['m_avg1'] + state['m_avg2']200

                state['v_avg'].lerp_(u.square(), 1 - group['beta_v'])

                u.div_(state['v_avg'].sqrt() + group['eps'])

                u.add_(p, alpha=group['weight_decay'])

                p.data.add_(u, alpha=-group['lr'])

                if group['lookahead_k'] > 0:
                    if state['step'] % group['lookahead_k'] == 0:
                        state['ema'].lerp_(p.data, 1 - group['lookahead_beta'])
                        p.data.copy_(state['ema'])

        return loss


def lilith_no_lookahead(*args, **kwargs):
    return Lilith(*args, **kwargs, lookahead_k=0)


if __name__ == "__main__":
    # python examples/viz_optimizers.py

    # Each optimizer has tweaked search space to produce better plots and
    # help to converge on better lr faster.
    optimizers = [(Lilith, -16, 0), (lilith_no_lookahead, -16, 0)]

    execute_experiments(optimizers, objective_rosenbrok, rosenbrock, plot_rosenbrok, (-2.0, 2.0), )



Now make necessary changes to this code to make animation:
from matplotlib.animation import FuncAnimation
import matplotlib.pyplot as plt
import numpy as np
import torch
from torch.optim import Adam, SGD, Adagrad, Adadelta, AdamW, radam, RMSprop
from tqdm import tqdm

# Define the Rosenbrock function
def rosenbrock(xy):
    x, y = xy
    return (1 - x) ** 2 + 100 * (y - x ** 2) ** 2

def run_optimization(xy_init, optimizer_class, n_iter, function, **optimizer_kwargs):
    xy_t = torch.tensor(xy_init, requires_grad=True)
    optimizer = optimizer_class([xy_t], **optimizer_kwargs)

    path = np.empty((n_iter + 1, 2))
    path[0, :] = xy_init

    for i in tqdm(range(1, n_iter + 1)):
        optimizer.zero_grad()
        loss = function(xy_t)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(xy_t, 1.0)
        optimizer.step()

        path[i, :] = xy_t.detach().numpy()

    return path

def create_animation(paths,
                     colors,
                     names,
                     function,
                     function_name,
                     figsize=(12, 12),
                     x_lim=(-2, 2),
                     y_lim=(-2, 2),
                     n_seconds=15):
    if not (len(paths) == len(colors) == len(names)):
        raise ValueError

    path_length = max(len(path) for path in paths)

    n_points = 1000
    x = np.linspace(*x_lim, n_points)
    y = np.linspace(*y_lim, n_points)
    X, Y = np.meshgrid(x, y)

    if function_name == "Rosenbrock":
        Z = rosenbrock([X, Y])
        minimum = (1.0, 1.0)

    fig, ax = plt.subplots(figsize=figsize)
    ax.contour(X, Y, Z, 90, cmap="jet")

    scatters = [ax.scatter(None,
                           None,
                           label=label,
                           c=c) for c, label in zip(colors, names)]

    ax.legend(prop={"size": 50})
    ax.plot(*minimum, "rD")

    def animate(i):
        for path, scatter in zip(paths, scatters):
            scatter.set_offsets(path[:i, :])

        ax.set_title(str(i))

    ms_per_frame = 1000 * n_seconds / path_length

    anim = FuncAnimation(fig, animate, frames=path_length, interval=ms_per_frame)

    return anim

if __name__ == "__main__":
    xy_init = (.3, .8)
    n_iter = 3000

    # path_adam_rosenbrock = run_optimization(xy_init, Adam, n_iter, rosenbrock)
    # path_sgd_rosenbrock = run_optimization(xy_init, SGD, n_iter, rosenbrock, lr=1e-3)

    freq = 10

    optimizers = [Adam, SM3,]#Adafactor,CAME,SM3]# Adagrad, RMSprop,
    paths_rosenbrock = [run_optimization(xy_init, optimizer, n_iter, rosenbrock,lr=3e-3)[::freq] for optimizer in optimizers]
    colors = ["green", "brown", ]#"purple", "brown", "magenta"]# "pink", "brown",
    names = [optimizer.__name__ for optimizer in optimizers]

    anim_rosenbrock = create_animation(paths_rosenbrock,
                                      colors,
                                      names,
                                      rosenbrock,
                                      "Rosenbrock",
                                      figsize=(15, 8),
                                      x_lim=(-2.1, 2.1),
                                      y_lim=(-2.1, 2.1),
                                      n_seconds=15)

    anim_rosenbrock.save("SM3_rosenbrock.gif")