From 4401ffe2e4bdaaff4331225a95f2beeaf1577156 Mon Sep 17 00:00:00 2001 From: Sean Perry Date: Thu, 1 Aug 2024 14:34:59 -0700 Subject: [PATCH] Update train.py --- pyha_analyzer/train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyha_analyzer/train.py b/pyha_analyzer/train.py index 5d4d26b..2249c96 100644 --- 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)