Fix forward() to return all multitask head outputs

app.py

@@ -88,7 +88,15 @@ class SpecialistDistilBERT(nn.Module):
         pooled  = outputs.last_hidden_state[:, 0, :]
         pooled  = self.dropout(pooled)
         shared  = self.shared(pooled)
-        return {"doc_logits": self.doc_classifier(shared)}
+        return {
+            "doc_logits":          self.doc_classifier(shared),
+            "expense_logits":      self.expense_classifier(shared),
+            "personal_logits":     self.personal_classifier(shared),
+            "business_logits":     self.business_classifier(shared),
+            "auto_renew_logits":   self.auto_renew_classifier(shared),
+            "renewal_logits":      self.renewal_period_classifier(shared),
+            "cancellation_logits": self.cancellation_classifier(shared),
+        }
 
 
 @asynccontextmanager