fix VQ for fp16

Vermeille · Oct 4, 2022 · e98572b · e98572b
1 parent 5fafd44
commit e98572b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchelie/nn/vq.py b/torchelie/nn/vq.py
@@ -66,7 +66,7 @@ def resample_dead(self, x):
             x_flat = x.view(-1, x.shape[-1])
             emb_weight = self.embedding.weight.data
             emb_weight[dead[:len(x_flat)]] = x_flat[torch.randperm(
-                len(x_flat))[:len(dead)]]
+                len(x_flat))[:len(dead)]].to(emb_weight.dtype)
             self.age[dead[:len(x_flat)]] = 0
 
             if torch.distributed.is_initialized():