Update app.py

app.py

@@ -215,13 +215,13 @@ def reconstruct(input_img, caption):
 
         update_scale(12)
 
-        real_cpu = real_image_initial_latents.detach().cpu()
-        inversed_cpu = [x.detach().cpu() for x in inversed_latents]
+        real_gpu = real_image_initial_latents.detach()
+        inversed_gpu = [x.detach() for x in inversed_latents]
 
         return image_np, caption, 12, [
             caption,
-            real_cpu,
-            inversed_cpu,
+            real_gpu,
+            inversed_gpu,
             weights
         ]