snowyday

from argparse import ArgumentParser

import torch
from torch import nn
import torch.nn.functional as F
from torch.optim import SGD
import torch.utils.data

# === ADDED FOR DISTRIBUTED >>>
import torch.distributed

snowyday

optimizer = optim.Eve(model.parameters())

optimizer.step(lambda: loss)

snowyday

import math
from torch.optim import Optimizer

class Eve(Optimizer):
    """Implements Eve (Adam with feedback) algorithm.
    
    It has been proposed in `Improving Stochastic Gradient Descent with Feedback, `_.
    
    Arguments:
        params (iterable): iterable of parameters to optimize or dicts defining