Update delta-iris/src/models/quantizer.py

delta-iris/src/models/quantizer.py

@@ -15,9 +15,9 @@ class Quantizer(nn.Module):
         self.pre_quant_proj = nn.Linear(input_dim, codebook_dim)
         self.post_quant_proj = nn.Linear(codebook_dim, input_dim)
         codebook = torch.empty(codebook_size, codebook_dim, requires_grad=False).uniform_(-1.0 / codebook_size, 1.0 / codebook_size)
-        self.num_codebook_updates = torch.tensor(0)
-        self.codebook = codebook
-        self.codewords_freqs = torch.ones(codebook_size).div(codebook_size)
+        self.register_buffer('num_codebook_updates', torch.tensor(0))
+        self.register_buffer('codebook', codebook)
+        self.register_buffer('codewords_freqs', torch.ones(codebook_size).div(codebook_size))
 
     def forward(self, z: torch.Tensor) -> dict:
         z = self.pre_quant_proj(z)