Upload model

rna_torsionbert_model.py

@@ -18,8 +18,8 @@ class RNATorsionBERTModel(PreTrainedModel):
             nn.Linear(self.dnabert_config.hidden_size, config.hidden_size),
             nn.GELU(),
             nn.Linear(config.hidden_size, config.num_classes),
-            nn.Softmax(dim=-1)
         )
+        self.activation = nn.Tanh()
 
     def init_model(self, k: int):
         model_name = f"zhihan1996/DNA_bert_{k}"
@@ -35,4 +35,5 @@ class RNATorsionBERTModel(PreTrainedModel):
     def forward(self, tensor):
         z = self.dnabert(**tensor).last_hidden_state
         output = self.regressor(z)
+        output = self.activation(output)
         return {"logits": output}