Update ddim.py: Fix the condition of devices.

GPU-Net · Aug 11, 2023 · ac00e4f · ac00e4f
1 parent 17a7ae2
commit ac00e4f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/audioldm2/latent_diffusion/models/ddim.py b/audioldm2/latent_diffusion/models/ddim.py
@@ -23,7 +23,7 @@ def __init__(self, model, schedule="linear", device=torch.device("cuda"), **kwar
  def register_buffer(self, name, attr):
  if type(attr) == torch.Tensor:
  if attr.device != self.device:
- if self.device == torch.device("mps") and attr.dtype==torch.float64:
+ if self.device == "mps" and attr.dtype==torch.float64:
  attr = attr.to(self.device, dtype=torch.float32)
  else:
  attr = attr.to(self.device)