Cuda uses now check for device

app.py

@@ -92,7 +92,6 @@ class ImageEditor(object):
 
         self.e4e_net = pSp(opts, self.device)
         self.e4e_net.eval()
-        self.e4e_net.cuda()
 
         self.shape_predictor = dlib.shape_predictor(
             model_paths["dlib"]
@@ -192,7 +191,7 @@ class ImageEditor(object):
 
     def run_on_batch(self, inputs):
         images, latents = self.e4e_net(
-            inputs.to("cuda").float(), randomize_noise=False, return_latents=True
+            inputs.to(self.device).float(), randomize_noise=False, return_latents=True
         )
         return images, latents
 

generate_videos.py

@@ -52,6 +52,8 @@ def project_code(latent_code, boundary, distance=3.0):
 
 def generate_frames(args, source_latent, g_ema_list, output_dir):
 
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
     alphas = np.linspace(0, 1, num=20)
     
     interpolate_func = interpolate_with_boundaries # default
@@ -84,7 +86,7 @@ def generate_frames(args, source_latent, g_ema_list, output_dir):
                 src_pars[k].data.copy_(mix_pars[segment_id][k] * (1 - mix_alpha) + mix_pars[segment_id + 1][k] * mix_alpha)
 
         if idx == 0 or segments or latent is not latents[idx - 1]:
-            w = torch.from_numpy(latent).float().cuda()
+            w = torch.from_numpy(latent).float().to(device)
 
             with torch.no_grad():
                 img, _ = g_ema([w], input_is_latent=True, truncation=1, randomize_noise=False)
@@ -205,7 +207,7 @@ def vid_to_gif(vid_path, output_dir, scale=256, fps=35):
 
 
 if __name__ == '__main__':
-    device = 'cuda'
+    device = "cuda" if torch.cuda.is_available() else "cpu"
 
     parser = argparse.ArgumentParser()