add attention model

Kapoorlabs-CAPED · Oct 9, 2024 · 2898c84 · 2898c84
1 parent 4aae5e8
commit 2898c84
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/kapoorlabs_lightning/pytorch_models.py b/src/kapoorlabs_lightning/pytorch_models.py
@@ -1585,6 +1585,8 @@ def forward(self, input):
     "TrackAsuraRotaryPositionalEncoding",
     "DenseNet",
     "MitosisNet",
+    "Attention",
+    "HybridAttentionDenseNet",
     "CloudAutoEncoder",
     "DGCNNEncoder",
     "FoldNetEncoder",