fix spwam

P3-SAM/demo/auto_mask.py

@@ -1253,7 +1253,8 @@ class AutoMask:
         self.model = P3SAM()
         self.model.load_state_dict(ckpt_path)
         self.model.eval()
-        self.model_parallel = torch.nn.DataParallel(self.model)
+        # self.model_parallel = torch.nn.DataParallel(self.model)
+        self.model_parallel = self.model
         self.model.cuda()
         self.model_parallel.cuda()
         self.point_num = point_num

P3-SAM/demo/auto_mask_no_postprocess.py

@@ -760,7 +760,8 @@ class AutoMask:
             self.model = P3SAM()
             self.model.load_state_dict(ckpt_path)
             self.model.eval()
-            self.model_parallel = torch.nn.DataParallel(self.model)
+            # self.model_parallel = torch.nn.DataParallel(self.model)
+            self.model_parallel = self.model
             self.model.cuda()
             self.model_parallel.cuda()
         self.point_num = point_num

XPart/partgen/bbox_estimator/auto_mask_api.py

@@ -1366,7 +1366,8 @@ class AutoMask:
             state_dict=torch.load(ckpt_path, weights_only=False, map_location="cpu")["state_dict"]
         )
         self.model.eval()
-        self.model_parallel = torch.nn.DataParallel(self.model)
+        # self.model_parallel = torch.nn.DataParallel(self.model)
+        self.model_parallel = self.model
         self.model.cuda()
         self.model_parallel.cuda()
         self.point_num = point_num

app.py

@@ -8,9 +8,9 @@ from pathlib import Path
 import torch
 import pytorch_lightning as pl
 import spaces
-import multiprocessing
-multiprocessing.set_start_method('spawn')
-print('using spawm')
+import torch.multiprocessing as mp
+mp.set_start_method('spawn')
+print('using torch spawm')
 
 
 sys.path.append('P3-SAM')