From 732257e8caa81abb5df7f50a8b63e464ba5e9479 Mon Sep 17 00:00:00 2001 From: Sean Perry Date: Thu, 1 Aug 2024 14:37:07 -0700 Subject: [PATCH] hotfix: save GPU space --- 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 68920b4..1c9a2ba 100644 --- a/pyha_analyzer/train.py +++ b/pyha_analyzer/train.py @@ -223,7 +223,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: PyhaDFDataset = data_loader.dataset # type: ignore cmap, smap = map_metric(log_pred, torch.cat(log_label), dataset.class_dist)