put costs on correct device

mittagessen · Sep 30, 2024 · cda9b18 · cda9b18
1 parent 1b577c2
commit cda9b18
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/kraken/lib/train/segmentation.py b/kraken/lib/train/segmentation.py
@@ -279,10 +279,11 @@ def validation_step(self, batch, batch_idx):
                     costs = np.sort(costs)[:len(y_curves[line_cls][0])]
                     penalty = np.full(diff, 8.0)
                     costs = np.concatenate([costs, penalty])
-                self.val_line_mean_dist.update(costs/8.0)
+                costs = costs/8.0
             # no line output
             else:
-                self.val_line_mean_dist.update(torch.ones(len(y_curves[line_cls][0])))
+                costs = torch.ones(len(y_curves[line_cls][0]))
+            self.val_line_mean_dist.update(costs.to(self.val_line_mean_dist.device))
 
     def on_validation_epoch_end(self):
         if not self.trainer.sanity_checking: