device

app.py

@@ -19,14 +19,14 @@ np.random.seed(0)
 
 def get_pil_im(im, resolution=64):
     im = imresize(im, (resolution, resolution))[:, :, :3]
-    im = th.Tensor(im).permute(2, 0, 1)[None, :, :, :].contiguous().cuda()
+    im = th.Tensor(im).permute(2, 0, 1)[None, :, :, :].contiguous()
     return im
 
 
 # generate image components and reconstruction
 def gen_image_and_components(model, gd, im, num_components=4, sample_method='ddim', batch_size=1, image_size=64, device='cuda', num_images=1):
     """Generate row of orig image, individual components, and reconstructed image"""
-    orig_img = get_pil_im(im, resolution=image_size)
+    orig_img = get_pil_im(im, resolution=image_size).to(device)
     latent = model.encode_latent(orig_img)
     model_kwargs = {'latent': latent}
 
@@ -74,7 +74,7 @@ def gen_image_and_components(model, gd, im, num_components=4, sample_method='ddi
 
 def decompose_image(im):
     sample_method = 'ddim'
-    result = gen_image_and_components(clevr_model, GD[sample_method], im, sample_method=sample_method, num_images=1)
+    result = gen_image_and_components(clevr_model, GD[sample_method], im, sample_method=sample_method, num_images=1, device=device)
     return result.permute(1, 2, 0).numpy()