fix: sort by cat_idx into embedding generator

dreamquark-ai · Jun 2, 2020 · 9ab3ad5 · 9ab3ad5
1 parent fd91028
commit 9ab3ad5
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/pytorch_tabnet/tab_network.py b/pytorch_tabnet/tab_network.py
@@ -424,6 +424,12 @@ def __init__(self, input_dim, cat_dims, cat_idxs, cat_emb_dim):
  self.post_embed_dim = int(input_dim + np.sum(self.cat_emb_dims) - len(self.cat_emb_dims))
 
  self.embeddings = torch.nn.ModuleList()
+
+ # Sort dims by cat_idx
+ sorted_idxs = np.argsort(cat_idxs)
+ cat_dims = [cat_dims[i] for i in sorted_idxs]
+ self.cat_emb_dims = [self.cat_emb_dims[i] for i in sorted_idxs]
+
  for cat_dim, emb_dim in zip(cat_dims, self.cat_emb_dims):
  self.embeddings.append(torch.nn.Embedding(cat_dim, emb_dim))