def _forward_auto_encoder(self, x, idx_scale):
    self.model.set_forward_ae_loss(True)
    x = self.model(x, idx_scale)
    self.model.set_forward_ae_loss(False)
    return x
def _forward_auto_encoder(self, x, idx_scale):
    self.model.set_forward_ae_loss(True)
    x = self.model(x, idx_scale)
    self.model.set_forward_ae_loss(False)
    return x