Merge pull request #133 from moritzschaefer/main

Fix missing `batch_first` attribute
bowang-lab · Dec 12, 2023 · db58fca · db58fca
2 parents 9f75191 + c32ce7d
commit db58fca
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/scgpt/model/model.py b/scgpt/model/model.py
@@ -641,6 +641,9 @@ def __init__(
             attention_dropout=dropout,
             **factory_kwargs,
         )
+        # Version compatibility workaround
+        if not hasattr(self.self_attn, "batch_first"):
+            self.self_attn.batch_first = batch_first
         # Implementation of Feedforward model
         self.linear1 = nn.Linear(d_model, dim_feedforward, **factory_kwargs)
         self.dropout = nn.Dropout(dropout)