def METHOD_NAME(self, model, inputs):
    """ Perform a training step on a batch of inputs.
    Subclass and override to inject custom behavior.
    Args:
        model (`TorchModel`): The model to train.
        inputs (`Dict[str, Union[torch.Tensor, Any]]`):
            The inputs and targets of the model.
            The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
            argument `labels`. Check your model's documentation for all accepted arguments.
    Return:
        `torch.Tensor`: The tensor with training loss on this batch.
    """
    # EvaluationHook will do evaluate and change mode to val, return to train mode
    # TODO: find more pretty way to change mode
    self.d_reg_every = self.cfg.train.get('d_reg_every', 16)
    self.g_reg_every = self.cfg.train.get('g_reg_every', 4)
    self.path_regularize = self.cfg.train.get('path_regularize', 2)
    self.r1 = self.cfg.train.get('r1', 10)
    train_outputs = dict()
    self._mode = ModeKeys.TRAIN
    # call model forward but not __call__ to skip postprocess
    if isinstance(inputs, Mapping):
        d_loss = model._train_forward_d(**inputs)
    else:
        d_loss = model._train_forward_d(inputs)
    train_outputs['d_loss'] = d_loss
    model.discriminator.zero_grad()
    d_loss.backward()
    self.optimizer_d.step()
    if self._iter % self.d_reg_every == 0:
        if isinstance(inputs, Mapping):
            r1_loss = model._train_forward_d_r1(**inputs)
        else:
            r1_loss = model._train_forward_d_r1(inputs)
        train_outputs['r1_loss'] = r1_loss
        model.discriminator.zero_grad()
        (self.r1 / 2 * r1_loss * self.d_reg_every).backward()
        self.optimizer_d.step()
    if isinstance(inputs, Mapping):
        g_loss = model._train_forward_g(**inputs)
    else:
        g_loss = model._train_forward_g(inputs)
    train_outputs['g_loss'] = g_loss
    model.generator.zero_grad()
    g_loss.backward()
    self.optimizer.step()
    path_loss = 0
    if self._iter % self.g_reg_every == 0:
        if isinstance(inputs, Mapping):
            path_loss = model._train_forward_g_path(**inputs)
        else:
            path_loss = model._train_forward_g_path(inputs)
        train_outputs['path_loss'] = path_loss
        model.generator.zero_grad()
        weighted_path_loss = self.path_regularize * self.g_reg_every * path_loss
        weighted_path_loss.backward()
        self.optimizer.step()
    model.accumulate()
    if not isinstance(train_outputs, dict):
        raise TypeError('"model.forward()" must return a dict')
    # add model output info to log
    if 'log_vars' not in train_outputs:
        default_keys_pattern = ['loss']
        match_keys = set([])
        for key_p in default_keys_pattern:
            match_keys.update(
                [key for key in train_outputs.keys() if key_p in key])
        log_vars = {}
        for key in match_keys:
            value = train_outputs.get(key, None)
            if value is not None:
                if dist.is_available() and dist.is_initialized():
                    value = value.data.clone()
                    dist.all_reduce(value.div_(dist.get_world_size()))
                log_vars.update({key: value.item()})
        self.log_buffer.update(log_vars)
    else:
        self.log_buffer.update(train_outputs['log_vars'])
    self.train_outputs = train_outputs