fix-478: add n_shared_decoder and n_indep_decoder

dreamquark-ai · Apr 25, 2023 · d44f3b5 · d44f3b5
1 parent 8d09d2f
commit d44f3b5
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/pytorch_tabnet/pretraining.py b/pytorch_tabnet/pretraining.py
@@ -188,6 +188,8 @@ def _set_network(self):
             cat_emb_dim=self.cat_emb_dim,
             n_independent=self.n_independent,
             n_shared=self.n_shared,
+            n_shared_decoder=self.n_shared_decoder,
+            n_indep_decoder=self.n_indep_decoder,
             epsilon=self.epsilon,
             virtual_batch_size=self.virtual_batch_size,
             momentum=self.momentum,