Пользовательский оптимизатор PyTorch не работает должным образом и не может получить доступ к градиентам

Я пытаюсь реализовать алгоритм эластичного усреднения стохастического градиентного спуска (EASGD) из статьи Глубокое обучение с эластичным усреднением SGD и столкнулся с некоторыми проблемами.

Я использую класс PyTorch torch.optim.Optimizer и ссылаюсь на официальную реализацию SGD и официальную реализацию Accelerated SGD, чтобы с чего-то начать.

Код, который у меня есть:

import torch.optim as optim


class EASGD(optim.Optimizer):
    def __init__(self, params, lr, tau, alpha=0.001):
        self.alpha = alpha

        if lr < 0.0:
            raise ValueError(f"Invalid learning rate {lr}.")

        defaults = dict(lr=lr, alpha=alpha, tau=tau)
        super(EASGD, self).__init__(params, defaults)

    def __setstate__(self, state):
        super(EASGD, self).__setstate__(state)

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

        for group in self.param_groups:
            tau = group['tau']

            for t, p in enumerate(group['params']):
                x_normal = p.clone()
                x_tilde = p.clone()

                if p.grad is None:
                    continue

                if t % tau == 0:
                    p = p - self.alpha * (x_normal - x_tilde)
                    x_tilde = x_tilde + self.alpha * (x_normal - x_tilde)

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

        return loss

Когда я запускаю этот код, я получаю следующую ошибку:

/home/user/github/test-repo/easgd.py:50: UserWarning: Доступ к атрибуту .grad тензора, который не является листовым тензором. Его атрибут .grad не будет заполнен во время autograd.backward(). Если вам действительно нужен градиент для нелистового тензора, используйте .retain_grad() на нелистовом тензоре. Если вы по ошибке получили доступ к нелистовому тензору, убедитесь, что вместо этого вы обращаетесь к листовому тензору. См. github.com/pytorch/pytorch/pull/30531 для получения дополнительной информации.

Чтение этого обсуждения PyTorch помогло понять, в чем разница между листовыми и нелистовыми переменными, но я не уверен, как мне исправить свой код, чтобы он работал правильно.

Любые советы о том, что делать или где искать приветствуются. Спасибо.