From 6056146701def08e82a969b6a3c32964bacec4f6 Mon Sep 17 00:00:00 2001 From: Benjamin Kiessling Date: Mon, 30 Sep 2024 12:54:13 +0200 Subject: [PATCH] dict-moving --- kraken/lib/train/segmentation.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/kraken/lib/train/segmentation.py b/kraken/lib/train/segmentation.py index 39b86b13..4dfa610e 100644 --- a/kraken/lib/train/segmentation.py +++ b/kraken/lib/train/segmentation.py @@ -267,7 +267,8 @@ def validation_step(self, batch, batch_idx): # cast pred/targets to float32 and move to CPU pred = pred.cpu().float() - y_curves = y_curves.cpu() + for k, v in y_curves.items(): + y_curves[k] = v.cpu() # vectorize and match lines for line_cls, line_idx in self.nn.user_metadata['class_mapping']['baselines'].items():