def train(self):
    self.loss.step()
    epoch = (self.optimizer.get_last_epoch() + 1)
    lr = self.optimizer.get_lr()
    self.ckp.write_log('[Epoch {}]\tLearning rate: {:.2e}'.format(epoch, Decimal(lr)))
    self.loss.start_log()
    self.model.train()
    (timer_data, timer_model) = (utility.timer(), utility.timer())
    self.loader_train.dataset.set_scale(0)
    for (batch, (lr, hr, _)) in enumerate(self.loader_train):
        (lr, hr) = self.prepare(lr, hr)
        if self.flag_ae_loss:
            (hr, hr_ae) = (hr[:, :self.args.n_colors, ...], hr[:, self.args.n_colors:, ...])
        else:
            hr_ae = None
        timer_data.hold()
        timer_model.tic()
        self.optimizer.zero_grad()
        if (self.args.precision == 'amp'):
            with torch.cuda.amp.autocast():
                sr = self.model(lr, 0)
                if self.flag_ae_loss:
                    sr_ae = self._forward_auto_encoder(hr_ae, 0)
                else:
                    sr_ae = None
                loss = self.loss(sr, hr, sr_ae, hr_ae)
            self.scaler.scale(loss).backward()
        else:
            sr = self.model(lr, 0)
            if self.flag_ae_loss:
                sr_ae = self._forward_auto_encoder(hr_ae, 0)
            else:
                sr_ae = None
            loss = self.loss(sr, hr, sr_ae, hr_ae)
            loss.backward()
        if (self.args.gclip > 0):
            utils.clip_grad_value_(self.model.parameters(), self.args.gclip)
        if (self.args.precision == 'amp'):
            self.scaler.step(self.optimizer)
            self.scaler.update()
        else:
            self.optimizer.step()
        timer_model.hold()
        if (((batch + 1) % self.args.print_every) == 0):
            self.ckp.write_log('[{}/{}]\t{}\t{:.1f}+{:.1f}s'.format(((batch + 1) * self.args.batch_size), len(self.loader_train.dataset), self.loss.display_loss(batch), timer_model.release(), timer_data.release()))
        timer_data.tic()
    self.loss.end_log(len(self.loader_train))
    self.error_last = self.loss.log[((- 1), (- 1))]
    self.optimizer.schedule()
def train(self):
    self.loss.step()
    epoch = (self.optimizer.get_last_epoch() + 1)
    lr = self.optimizer.get_lr()
    self.ckp.write_log('[Epoch {}]\tLearning rate: {:.2e}'.format(epoch, Decimal(lr)))
    self.loss.start_log()
    self.model.train()
    (timer_data, timer_model) = (utility.timer(), utility.timer())
    self.loader_train.dataset.set_scale(0)
    for (batch, (lr, hr, _)) in enumerate(self.loader_train):
        (lr, hr) = self.prepare(lr, hr)
        if self.flag_ae_loss:
            (hr, hr_ae) = (hr[:, :self.args.n_colors, ...], hr[:, self.args.n_colors:, ...])
        else:
            hr_ae = None
        timer_data.hold()
        timer_model.tic()
        self.optimizer.zero_grad()
        if (self.args.precision == 'amp'):
            with torch.cuda.amp.autocast():
                sr = self.model(lr, 0)
                if self.flag_ae_loss:
                    sr_ae = self._forward_auto_encoder(hr_ae, 0)
                else:
                    sr_ae = None
                loss = self.loss(sr, hr, sr_ae, hr_ae)
            self.scaler.scale(loss).backward()
        else:
            sr = self.model(lr, 0)
            if self.flag_ae_loss:
                sr_ae = self._forward_auto_encoder(hr_ae, 0)
            else:
                sr_ae = None
            loss = self.loss(sr, hr, sr_ae, hr_ae)
            loss.backward()
        if (self.args.gclip > 0):
            utils.clip_grad_value_(self.model.parameters(), self.args.gclip)
        if (self.args.precision == 'amp'):
            self.scaler.step(self.optimizer)
            self.scaler.update()
        else:
            self.optimizer.step()
        timer_model.hold()
        if (((batch + 1) % self.args.print_every) == 0):
            self.ckp.write_log('[{}/{}]\t{}\t{:.1f}+{:.1f}s'.format(((batch + 1) * self.args.batch_size), len(self.loader_train.dataset), self.loss.display_loss(batch), timer_model.release(), timer_data.release()))
        timer_data.tic()
    self.loss.end_log(len(self.loader_train))
    self.error_last = self.loss.log[((- 1), (- 1))]
    self.optimizer.schedule()