fix: n_a not being used

dreamquark-ai · Jan 2, 2021 · 7ae20c9 · 7ae20c9
1 parent 61c294a
commit 7ae20c9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pytorch_tabnet/tab_network.py b/pytorch_tabnet/tab_network.py
@@ -331,7 +331,7 @@ def __init__(
             input_dim=self.post_embed_dim,
             output_dim=self.post_embed_dim,
             n_d=n_d,
-            n_a=n_d,
+            n_a=n_a,
             n_steps=n_steps,
             gamma=gamma,
             n_independent=n_independent,
@@ -442,7 +442,7 @@ def __init__(
             input_dim=input_dim,
             output_dim=output_dim,
             n_d=n_d,
-            n_a=n_d,
+            n_a=n_a,
             n_steps=n_steps,
             gamma=gamma,
             n_independent=n_independent,