update

projects/video_diffusion_sr/infer.py

@@ -71,7 +71,7 @@ class VideoDiffusionInfer():
 
     @log_on_entry
     @log_runtime
-    def configure_dit_model(self, device="cpu", checkpoint=None):
+    def configure_dit_model(self, device="cuda", checkpoint=None):
         # Load dit checkpoint.
         # For fast init & resume,
         #   when training from scratch, rank0 init DiT on cpu, then sync to other ranks with FSDP.
@@ -83,7 +83,7 @@ class VideoDiffusionInfer():
         self.dit.set_gradient_checkpointing(self.config.dit.gradient_checkpoint)
 
         if checkpoint:
-            state = torch.load(checkpoint, map_location="cpu", mmap=True)
+            state = torch.load(checkpoint, map_location=self.device, mmap=True)
             loading_info = self.dit.load_state_dict(state, strict=True, assign=True)
             print(f"Loading pretrained ckpt from {checkpoint}")
             print(f"Loading info: {loading_info}")