Models now loaded on appropriate device

app.py

@@ -69,7 +69,7 @@ class ImageEditor(object):
                 model_size, latent_size, n_mlp, channel_multiplier=channel_mult
             ).to(self.device)
 
-            checkpoint = torch.load(model_paths[model])
+            checkpoint = torch.load(model_paths[model], map_location=self.device)
 
             g_ema.load_state_dict(checkpoint['g_ema'])