1 files changed, 10 insertions, 3 deletions

diff --git a/models/sparse.py b/models/sparse.py
index d706db5..bcb2897 100644
--- a/models/sparse.py
+++ b/models/sparse.py
@@ -5,22 +5,29 @@ import torch.nn as nn
 
 
 class PseudoSparseEmbedding(nn.Module):
-    def __init__(self, embedding_dim: int, device=None, dtype=torch.float32):
+    def __init__(self, embedding_dim: int, dropout_p: float = 0.0, device=None, dtype=torch.float32):
         super().__init__()
 
         self.embedding_dim = embedding_dim
         self.dtype = dtype
         self.params = nn.ParameterList()
+
+        if dropout_p > 0.0:
+            self.dropout = nn.Dropout(p=dropout_p)
+        else:
+            self.dropout = lambda x: x
+
         self.register_buffer('mapping', torch.zeros(0, device=device, dtype=torch.long))
 
     def forward(self, input_ids: torch.LongTensor):
-        ids = self.mapping[input_ids.to(self.mapping.device)]
+        input_ids = input_ids.to(self.mapping.device)
+        ids = self.mapping[input_ids]
         mask = ~(ids == -1)
 
         if torch.all(~mask):
             embs = None
         else:
-            embs = torch.stack([self.params[id] for id in ids[mask]])
+            embs = self.dropout(torch.stack([self.params[id] for id in ids[mask]]))
 
         return embs, mask