delete self.paper attribute from mambatab

basf · Jul 17, 2024 · 864d4f0 · 864d4f0
1 parent 96e9449
commit 864d4f0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mambular/base_models/mambatab.py b/mambular/base_models/mambatab.py
@@ -99,7 +99,7 @@ def forward(self, num_features, cat_features):
 
         x = self.norm_f(x)
         x = self.embedding_activation(x)
-        if self.paper:
+        if self.axis == 1:
             x = x.squeeze(1)
         else:
             x = x.squeeze(0)