backed out change from PR invoke-ai#44 that was causing ddim sampler …

…to fail with the message 'sqrt _vml_cpu not implemented for 'Half'
cognidesign · Aug 24, 2022 · ae7ee93 · ae7ee93
1 parent 7ecd6fd
commit ae7ee93
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/ldm/models/diffusion/ddim.py b/ldm/models/diffusion/ddim.py
@@ -17,6 +17,9 @@ def __init__(self, model, schedule="linear", **kwargs):
         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"))
         setattr(self, name, attr)
 
     def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):