dict-moving

mittagessen · Sep 30, 2024 · 6056146 · 6056146
1 parent f2199c8
commit 6056146
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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():