def _add_loss_to_state(self, state: RunnerState, loss):
if self.loss_key is None:
if state.loss is not None:
if isinstance(state.loss, list):
state.loss.append(loss)
else:
state.loss = [state.loss, loss]
else:
state.loss = loss
else:
if state.loss is not None:
assert isinstance(state.loss, dict)
state.loss[self.loss_key] = loss
else:
state.loss = {self.loss_key: loss}
def _add_loss_to_state(loss_key: Optional[str], state: RunnerState,
loss: torch.Tensor):
if loss_key is None:
if state.loss is not None:
if isinstance(state.loss, list):
state.loss.append(loss)
else:
state.loss = [state.loss, loss]
else:
state.loss = loss
else:
if state.loss is not None:
assert isinstance(state.loss, dict)
state.loss[loss_key] = loss
else:
state.loss = {loss_key: loss}