Attempting to solve pickle issue

demo.py

@@ -23,8 +23,10 @@ from mast3r.utils.misc import hash_md5
 import main
 import utils.export as export
 
-@spaces.GPU(duration=10)
-def get_reconstructed_scene(outdir, model, device, silent, image_size, ios_mode, filelist):
+@spaces.GPU(duration=15)
+def get_reconstructed_scene(outdir, model, silent, image_size, ios_mode, filelist):
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
     assert len(filelist) == 1 or len(filelist) == 2, "Please provide one or two images"
     if ios_mode:
@@ -37,6 +39,7 @@ def get_reconstructed_scene(outdir, model, device, silent, image_size, ios_mode,
         img['img'] = img['img'].to(device)
         img['original_img'] = img['original_img'].to(device)
         img['true_shape'] = torch.from_numpy(img['true_shape'])
+    model = model.to(device)
 
     output = model(imgs[0], imgs[1])
 
@@ -50,13 +53,11 @@ if __name__ == '__main__':
     image_size = 512
     silent = False
     ios_mode = True
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
     model_name = "brandonsmart/splatt3r_v1.0"
     filename = "epoch=19-step=1200.ckpt"
     weights_path = hf_hub_download(repo_id=model_name, filename=filename)
     model = main.MAST3RGaussians.load_from_checkpoint(weights_path, 'cpu')
-    model = model.to(device)
     chkpt_tag = hash_md5(weights_path)
 
     # Define example inputs and their corresponding precalculated outputs
@@ -88,7 +89,7 @@ if __name__ == '__main__':
         cache_path = os.path.join(tmpdirname, chkpt_tag)
         os.makedirs(cache_path, exist_ok=True)
 
-        recon_fun = functools.partial(get_reconstructed_scene, tmpdirname, model, device, silent, image_size, ios_mode)
+        recon_fun = functools.partial(get_reconstructed_scene, tmpdirname, model, silent, image_size, ios_mode)
 
         if not ios_mode:
             for i in range(len(examples)):