def METHOD_NAME(self, batch):
    self.msg = "Rank: {}, ".format(dist.get_rank())
    losses_dict = {}
    # parse batch
    self.model.train()
    self.optimizer.clear_grad()
    wav, y, mel = batch
    y_hat = self.model(wav, mel)
    if self.mode == 'RAW':
        y_hat = y_hat.transpose([0, 2, 1]).unsqueeze(-1)
    elif self.mode == 'MOL':
        y_hat = paddle.cast(y, dtype='float32')
    y = y.unsqueeze(-1)
    loss = self.criterion(y_hat, y)
    loss.backward()
    grad_norm = float(
        calculate_grad_norm(self.model.parameters(), norm_type=2))
    self.optimizer.step()
    report("train/loss", float(loss))
    report("train/grad_norm", float(grad_norm))
    losses_dict["loss"] = float(loss)
    losses_dict["grad_norm"] = float(grad_norm)
    self.msg += ', '.join('{}: {:>.6f}'.format(k, v)
                          for k, v in losses_dict.items())
    global ITERATION
    ITERATION = self.state.iteration + 1