output_activation = hydra.utils.instantiate(
config.output_activation, _partial_=True
)
preds = output_activation(logits)