mend

mims-harvard · Jan 10, 2025 · 80806fa · 80806fa
1 parent d22f033
commit 80806fa
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tdc/model_server/models/scgpt.py b/tdc/model_server/models/scgpt.py
@@ -39,6 +39,7 @@ def __init__(
         self.cell_emb_style = cell_emb_style
         self.norm_scheme = norm_scheme
         self.explicit_zero_prob = explicit_zero_prob
+        self.use_flash_attention = self.use_fast_transformer
         super().__init__(pad_token_id=pad_token_id, **kwargs)
 
 class ExprDecoder(nn.Module):