Update CatVTON/model/DensePose/__init__.py

CatVTON/model/DensePose/__init__.py

@@ -38,11 +38,14 @@ class DensePose:
         self.min_score = 0.8
 
         self.cfg = self.setup_config()
+        self.cfg.MODEL.DEVICE = self.device  # Force device before creating predictor
         self.predictor = DefaultPredictor(self.cfg)
-        self.predictor.model.to(self.device)
 
     def setup_config(self):
-        opts = ["MODEL.ROI_HEADS.SCORE_THRESH_TEST", str(self.min_score)]
+        opts = [
+            "MODEL.ROI_HEADS.SCORE_THRESH_TEST", str(self.min_score),
+            "MODEL.DEVICE", self.device  # Add this line
+        ]
         cfg = get_cfg()
         add_densepose_config(cfg)
         cfg.merge_from_file(self.config_path)