net.to(self.device) update

app_function.py

@@ -260,6 +260,8 @@ class YOND_Backend:
                 raw_vst = (raw_vst - lower) / (upper - lower)
 
                 ################# 准备去噪 #################
+                self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+                self.yond.net = self.yond.net.to(self.device)
                 raw_vst = torch.from_numpy(raw_vst).float().to(self.device).permute(2,0,1)[None,]
                 if 'guided' in self.yond.arch:
                     sigma_corr = 1.03
@@ -333,6 +335,8 @@ class YOND_Backend:
             raw_vst = (raw_vst - lower) / (upper - lower)
             
             ################# 准备去噪 #################
+            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            self.yond.net = self.yond.net.to(self.device)
             raw_vst = torch.from_numpy(raw_vst).float().to(self.device).permute(2,0,1)[None,]
             if 'guided' in self.yond.arch:
                 t = torch.tensor(nsr*self.p['sigsnr'], dtype=raw_vst.dtype, device=self.device)
@@ -493,5 +497,4 @@ class YOND_anytest():
         # 模型加载
         self.net = globals()[self.arch['name']](self.arch)
         model = torch.load(model_path, map_location='cpu')
-        self.net = load_weights(self.net, model, by_name=False)
-        self.net = self.net.to(self.device)
+        self.net = load_weights(self.net, model, by_name=False)