Fixed #60 and #61. Placing dropout in the right position and applying…

… FullEmbeddingDropout only in training

.gitignore

@@ -11,7 +11,7 @@ __pycache__*
 Untitled*.ipynb
 
 # data related dirs
-# data/
+tmp_data/
 model_weights/
 tmp_dir/
 weights/

pytorch_widedeep/models/transformers/_embeddings_layers.py

@@ -18,10 +18,13 @@ def __init__(self, dropout: float):
         self.dropout = dropout
 
     def forward(self, X: Tensor) -> Tensor:
-        mask = X.new().resize_((X.size(1), 1)).bernoulli_(1 - self.dropout).expand_as(
-            X
-        ) / (1 - self.dropout)
-        return mask * X
+        if self.training:
+            mask = X.new().resize_((X.size(1), 1)).bernoulli_(1 - self.dropout).expand_as(
+                X
+            ) / (1 - self.dropout)
+            return mask * X
+        else:
+            return X
 
 
 DropoutLayers = Union[nn.Dropout, FullEmbeddingDropout]
@@ -167,12 +170,11 @@ def forward(self, X: Tensor) -> Tensor:
             x = torch.cat(cat_embed, 1)
         else:
             x = self.embed(X[:, self.cat_idx].long())
-            x = self.dropout(x)
 
         if self.bias is not None:
             x = x + self.bias.unsqueeze(0)
 
-        return x
+        return self.dropout(x)
 
 
 class CatAndContEmbeddings(nn.Module):