update

gradio_demo/app.py

@@ -126,7 +126,8 @@ def instantir_restore(
     print("Loading checkpoint...")
     aggregator_state_dict = torch.load(
         f"{instantir_path}/aggregator.pt",
-        map_location="cpu"
+        map_location = device if not cpu_offload else "cpu"
+#        map_location="cpu"
     )
     pipe.aggregator.load_state_dict(aggregator_state_dict, strict=True)
     if not cpu_offload: