update for BatchQC.

STOmics · Oct 17, 2023 · 8a58b88 · 8a58b88
1 parent f4c4d58
commit 8a58b88
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/stereo/algorithm/batch_qc/module/classifier.py b/stereo/algorithm/batch_qc/module/classifier.py
@@ -115,9 +115,9 @@ def train(self, max_epochs=500, save_path=None):
             epoch_loss = []
             for idx, data in enumerate(self.train_loader):
                 x, y = data
-                x = x.to(self.device, non_blocking=True)
+                x = x.float().to(self.device, non_blocking=True)
                 y = y.long().to(self.device)
-                if torch.cuda.is_available():
+                if self.device.type == 'cuda' and torch.cuda.is_available():
                     with torch.cuda.amp.autocast():
                         y_hat = self.model(x)
                         loss = self.loss_fn(y_hat, y)
@@ -157,7 +157,7 @@ def validation(self, pt_path):
         batch_acc = []
         for idx, data in enumerate(self.train_loader):
             x, y = data
-            x = x.to(self.device)
+            x = x.float().to(self.device)
             y = y.long().to(self.device)
             with torch.no_grad():
                 y_hat = self.model(x)
@@ -179,7 +179,7 @@ def test(self, pt_path):
         batch_acc = []
         for idx, data in enumerate(self.test_loader):
             x, y = data
-            x = x.to(self.device)
+            x = x.float().to(self.device)
             y = y.long().to(self.device)
             with torch.no_grad():
                 y_hat = self.model(x)