def METHOD_NAME(self):
    if self.opt and not self.SKIP_ZERO_GRAD:
        self.opt.zero_grad()
    for data, target in zip(self.example_inputs, self.example_outputs):
        with self.amp_context():
            pred = self.model(data)
        self.loss_fn(pred, target).backward()
        if self.opt:
            self.opt.step()