output_activation:
_target_: "torch.softmax"
dim: 1