Update device.py

device.py

@@ -14,8 +14,8 @@ from torchvision import transforms as tfms
 from tqdm.auto import tqdm
 from transformers import CLIPTextModel, CLIPTokenizer, logging
 import os
-#torch_device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
-torch_device = "cpu" 
+torch_device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+#torch_device = "cpu" 
 
 
 # Load the autoencoder model which will be used to decode the latents into image space.