Spaces:

Purple11
/

Grounded-Diffusion

Runtime error

Purple11 commited on Feb 17, 2023

Commit

0b162ec

1 Parent(s): 510b080

Update ldm/models/diffusion/ddim.py

Files changed (1) hide show

ldm/models/diffusion/ddim.py CHANGED Viewed

@@ -23,7 +23,7 @@ class DDIMSampler(object):
     def register_buffer(self, name, attr):
         if type(attr) == torch.Tensor:
             if attr.device != torch.device("cuda"):
-                attr = attr.to(torch.device("cuda"))
         setattr(self, name, attr)
     def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):

     def register_buffer(self, name, attr):
         if type(attr) == torch.Tensor:
             if attr.device != torch.device("cuda"):
+                attr = attr.to(torch.device("cuda")) if torch.cuda.is_available() else attr.to(torch.device("cpu"))
         setattr(self, name, attr)
     def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):