.

src/1models/LGATr/lgatr.py

@@ -1,4 +1,9 @@
-from lgatr import GATr, SelfAttentionConfig, MLPConfig
+try:
+    from lgatr import LGATr, SelfAttentionConfig, MLPConfig
+except:
+    print("Failed importing, trying importing GATr")
+    from lgatr import SelfAttentionConfig, MLPConfig
+    from lgatr import GATr as LGATr
 from lgatr.interface import embed_vector, extract_scalar, embed_spurions, extract_vector
 import torch
 import torch.nn as nn
@@ -19,7 +24,7 @@ class LGATrModel(torch.nn.Module):
         self.n_scalars_out = n_scalars_out
         self.obj_score = obj_score
         self.global_features_copy = global_featuers_copy
-        self.gatr = GATr(
+        self.gatr = LGATr(
             in_mv_channels=3,
             out_mv_channels=1,
             hidden_mv_channels=hidden_mv_channels,