ysharma HF staff commited on
Commit
c07d35e
1 Parent(s): a648d42
Files changed (1) hide show
  1. ddim.py +2 -1
ddim.py CHANGED
@@ -18,7 +18,8 @@ class DDIMSampler(object):
18
  def register_buffer(self, name, attr):
19
  if type(attr) == torch.Tensor:
20
  if attr.device != torch.device("cuda"):
21
- attr = attr.to(torch.device("cuda"))
 
22
  setattr(self, name, attr)
23
 
24
  def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
 
18
  def register_buffer(self, name, attr):
19
  if type(attr) == torch.Tensor:
20
  if attr.device != torch.device("cuda"):
21
+ continue
22
+ #attr = attr.to(torch.device("cuda"))
23
  setattr(self, name, attr)
24
 
25
  def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):