Update train.py

UCSD-E4E · Aug 1, 2024 · 4401ffe · 4401ffe
1 parent 3cf111d
commit 4401ffe
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pyha_analyzer/train.py b/pyha_analyzer/train.py
@@ -224,7 +224,7 @@ def valid(model: Any,
 
 
     # softmax predictions
-    log_pred = F.sigmoid(torch.cat(log_pred)).to(cfg.device)
+    log_pred = F.sigmoid(torch.cat(log_pred)).cpu() #.to(cfg.device)
 
     dataset = data_loader.dataset # type: ignore
     cmap, smap = map_metric(log_pred, torch.cat(log_label), dataset.class_dist)