Fix check for making sure vae is on the right device.

toolkit/models/base_model.py

@@ -1084,7 +1084,7 @@ class BaseModel:
 
         latent_list = []
         # Move to vae to device if on cpu
-        if self.vae.device == 'cpu':
+        if self.vae.device == torch.device("cpu"):
             self.vae.to(device)
         self.vae.eval()
         self.vae.requires_grad_(False)
@@ -1127,7 +1127,7 @@ class BaseModel:
             dtype = self.torch_dtype
 
         # Move to vae to device if on cpu
-        if self.vae.device == 'cpu':
+        if self.vae.device == torch.device('cpu'):
             self.vae.to(self.device)
         latents = latents.to(device, dtype=dtype)
         latents = (