From 9f5f70cb2a8919cb07821f264bf0fd75bfa10584 Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Sat, 8 Apr 2023 17:38:49 +0200
Subject: Update

---
 models/sparse.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'models/sparse.py')

diff --git a/models/sparse.py b/models/sparse.py
index 8910316..d706db5 100644
--- a/models/sparse.py
+++ b/models/sparse.py
@@ -11,7 +11,7 @@ class PseudoSparseEmbedding(nn.Module):
         self.embedding_dim = embedding_dim
         self.dtype = dtype
         self.params = nn.ParameterList()
-        self.mapping = torch.zeros(0, device=device, dtype=torch.long)
+        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)]
-- 
cgit v1.2.3-54-g00ecf