From 3b4c0431aa6ca68a7286532639ffaf6bd765a1d9 Mon Sep 17 00:00:00 2001 From: baskrahmer Date: Mon, 29 Apr 2024 17:11:53 +0200 Subject: [PATCH] Add SoftPlus to Hybrid head --- src/model.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/model.py b/src/model.py index 30b508d..f780467 100644 --- a/src/model.py +++ b/src/model.py @@ -69,12 +69,13 @@ def __init__(self, hidden_dim: int, num_classes: int, alpha: float): self.regression_linear = nn.Linear(in_features=num_classes, out_features=1) self.regression_loss = nn.MSELoss() self.classification_loss = nn.CrossEntropyLoss() + self.activation = nn.Softplus() self.alpha = alpha def __call__(self, activations, classes, regressands, **kwargs) -> dict: logits = self.classification_linear(activations) classification_loss = self.classification_loss(logits, classes) - predicted_values = self.regression_linear(logits) + predicted_values = self.activation(self.regression_linear(logits)) regression_loss = self.regression_loss(predicted_values, regressands) return { "predicted_values": predicted_values,