update

projects/video_diffusion_sr/infer.py

@@ -75,13 +75,13 @@ class VideoDiffusionInfer():
         # For fast init & resume,
         #   when training from scratch, rank0 init DiT on cpu, then sync to other ranks with FSDP.
         #   otherwise, all ranks init DiT on meta device, then load_state_dict with assign=True.
-        if self.config.dit.get("init_with_meta_device", False):
-            init_device = "cpu" if get_global_rank() == 0 and checkpoint is None else "meta"
-        else:
-            init_device = "cpu"
+        # if self.config.dit.get("init_with_meta_device", False):
+        #     init_device = "cpu" if get_global_rank() == 0 and checkpoint is None else "meta"
+        # else:
+        #     init_device = "cpu"
 
         # Create dit model.
-        with torch.device(init_device):
+        with torch.device("cpu"):
             self.dit = create_object(self.config.dit.model)
         self.dit.set_gradient_checkpointing(self.config.dit.gradient_checkpoint)
 
@@ -92,8 +92,8 @@ class VideoDiffusionInfer():
             print(f"Loading info: {loading_info}")
             self.dit = meta_non_persistent_buffer_init_fn(self.dit)
 
-        if device in [get_device(), "cuda"]:
-            self.dit.to(get_device())
+        # if device in [get_device(), "cuda"]:
+        self.dit.to("cuda")
 
         # Print model size.
         num_params = sum(p.numel() for p in self.dit.parameters() if p.requires_grad)
@@ -106,11 +106,11 @@ class VideoDiffusionInfer():
         dtype = getattr(torch, self.config.vae.dtype)
         self.vae = create_object(self.config.vae.model)
         self.vae.requires_grad_(False).eval()
-        self.vae.to(device=get_device(), dtype=dtype)
+        self.vae.to(device="cuda", dtype=dtype)
 
         # Load vae checkpoint.
         state = torch.load(
-            self.config.vae.checkpoint, map_location=get_device(), mmap=True
+            self.config.vae.checkpoint, map_location="cuda", mmap=True
         )
         self.vae.load_state_dict(state)