From ac5f2b33dcd156929f1ccbd966d095fdfc6dc251 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 13 Mar 2025 05:38:21 -0700 Subject: [PATCH] Update _utils.py --- unsloth/models/_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 680402a178..1a8fff9ada 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -986,7 +986,9 @@ def _unsloth_pre_compute_loss(self, model, inputs, *args, **kwargs): "Read more on gradient accumulation issues here: https://unsloth.ai/blog/gradient" ) pass - return self._old_compute_loss(model, inputs, *args, **kwargs) + with torch.autocast(device_type = "cuda", dtype = torch.float32): + outputs = self._old_compute_loss(model, inputs, *args, **kwargs) + return outputs pass