Move loss metric to same device as inputs (#3522)

ludwig-ai · Aug 11, 2023 · 007c261 · 007c261
1 parent c545285
commit 007c261
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/ludwig/models/llm.py b/ludwig/models/llm.py
@@ -484,7 +484,11 @@ def update_metrics(self, targets, predictions):
                 continue
             of_obj.update_metrics(targets[of_name], predictions[of_name])
 
+        # HACK (Tim): get the device of the targets to transfer self.eval_loss_metric to the same device
+        target_device = list(targets.values())[0].device
+
         eval_loss, additional_losses = self.eval_loss(targets, predictions)
+        self.eval_loss_metric = self.eval_loss_metric.to(target_device)
         self.eval_loss_metric.update(eval_loss)
         self.eval_additional_losses_metrics.update(additional_losses)