Aore commited on
Commit
dc96674
1 Parent(s): fe8c9b6

Update cldm/ddim_hacked.py

Browse files
Files changed (1) hide show
  1. cldm/ddim_hacked.py +2 -2
cldm/ddim_hacked.py CHANGED
@@ -16,8 +16,8 @@ class DDIMSampler(object):
16
 
17
  def register_buffer(self, name, attr):
18
  if type(attr) == torch.Tensor:
19
- if attr.device != torch.device("cuda"):
20
- attr = attr.to(torch.device("cuda"))
21
  setattr(self, name, attr)
22
 
23
  def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
 
16
 
17
  def register_buffer(self, name, attr):
18
  if type(attr) == torch.Tensor:
19
+ if attr.device != torch.device("cpu"):
20
+ attr = attr.to(torch.device("cpu"))
21
  setattr(self, name, attr)
22
 
23
  def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):