Merge branch 'feature/mamba' of github.com:VectorInstitute/odyssey in…

…to feature/mamba
VectorInstitute · May 1, 2024 · 8603ed6 · 8603ed6
2 parents 9e70a4b + f2910c9
commit 8603ed6
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/odyssey/models/cehr_mamba/mamba_utils.py b/odyssey/models/cehr_mamba/mamba_utils.py
@@ -1,5 +1,6 @@
 """Utilities following HuggingFace style for Mamba models."""
 
+from dataclasses import dataclass
 from typing import Optional, Tuple, Union
 
 import torch
@@ -19,7 +20,6 @@
     add_start_docstrings_to_model_forward,
     replace_return_docstrings,
 )
-from dataclasses import dataclass
 
 
 _CONFIG_FOR_DOC = "MambaConfig"
@@ -105,7 +105,8 @@ def forward(
             config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
             `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
 
-        Returns:
+        Returns
+        -------
         """
         sequence_outputs = self.backbone(
             input_ids,