fix: remove torch warnings (index should be bool)

dreamquark-ai · Nov 8, 2019 · f5817cf · f5817cf
1 parent 18ec79b
commit f5817cf
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytorch_tabnet/tab_network.py b/pytorch_tabnet/tab_network.py
@@ -117,7 +117,7 @@ def __init__(self, input_dim, output_dim, n_d, n_a,
             self.embeddings.append(torch.nn.Embedding(cat_dim, emb_dim))
 
         # record continuous indices
-        self.continuous_idx = torch.ones(self.input_dim, dtype=torch.uint8)
+        self.continuous_idx = torch.ones(self.input_dim, dtype=torch.bool)
         self.continuous_idx[self.cat_idxs] = 0
         self.post_embed_dim = self.input_dim + (cat_emb_dim - 1)*len(self.cat_idxs)
         self.initial_bn = BatchNorm1d(self.post_embed_dim, momentum=0.01)