cuda

ldm/guided_diffusion/h_posterior.py

@@ -60,10 +60,11 @@ class HPosterior(object):
         mean = (var) * ( (alpha_s_t_star/scale_s_t_star**2) * z_t_star +  (alpha_t_s/scale_t_s**2) * z_t )
         return mean, torch.sqrt(var)
     
+
     def register_buffer(self, name, attr):
-        if type(attr) == torch.Tensor:
-            if attr.device != torch.device("cuda"):
-                attr = attr.to(torch.device("cuda"))
+        if isinstance(attr, torch.Tensor):
+            if not attr.is_cuda:
+                attr = attr.cuda()
         setattr(self, name, attr)
 
     def get_error(self,x,t,c, unconditional_conditioning, unconditional_guidance_scale):

ldm/models/diffusion/ddim.py

@@ -16,9 +16,8 @@ class DDIMSampler(object):
         self.schedule = schedule
 
     def register_buffer(self, name, attr):
-        if type(attr) == torch.Tensor:
-            if attr.device != torch.device("cuda"):
-                attr = attr.to(torch.device("cuda"))
+        if isinstance(attr, torch.Tensor) and not attr.is_cuda:
+            attr = attr.cuda()
         setattr(self, name, attr)
 
     def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):

ldm/models/diffusion/plms.py

@@ -16,9 +16,8 @@ class PLMSSampler(object):
         self.schedule = schedule
 
     def register_buffer(self, name, attr):
-        if type(attr) == torch.Tensor:
-            if attr.device != torch.device("cuda"):
-                attr = attr.to(torch.device("cuda"))
+        if isinstance(attr, torch.Tensor) and not attr.is_cuda:
+            attr = attr.cuda()
         setattr(self, name, attr)
 
     def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):