MeYourHint commited on
Commit
b0132ea
·
1 Parent(s): 9d4e55b
Files changed (1) hide show
  1. models/vq/quantizer.py +1 -1
models/vq/quantizer.py CHANGED
@@ -44,7 +44,7 @@ class QuantizeEMAReset(nn.Module):
44
  self.init = False
45
  self.code_sum = None
46
  self.code_count = None
47
- self.register_buffer('codebook', torch.zeros(self.nb_code, self.code_dim, requires_grad=False).cuda())
48
 
49
  def _tile(self, x):
50
  nb_code_x, code_dim = x.shape
 
44
  self.init = False
45
  self.code_sum = None
46
  self.code_count = None
47
+ self.register_buffer('codebook', torch.zeros(self.nb_code, self.code_dim, requires_grad=False))
48
 
49
  def _tile(self, x):
50
  nb_code_x, code_dim = x.shape