ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update _utils.py

Browse source
This commit is contained in:
Daniel Han 2025-03-16 23:37:46 -07:00
commit 9ff6fde1f0
1 changed files with 7 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

13
unsloth/models/_utils.py
View file

@ -1017,12 +1017,13 @@ def _unsloth_pre_compute_loss(self, model, inputs, *args, **kwargs):
)
pass
if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "0":
autocaster = contextlib.nullcontext()
else:
autocaster = torch.autocast(device_type = "cuda", dtype = torch.float32)
with autocaster:
outputs = self._old_compute_loss(model, inputs, *args, **kwargs)
# if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "0":
# autocaster = contextlib.nullcontext()
# else:
# autocaster = torch.autocast(device_type = "cuda", dtype = torch.float32)
# with autocaster:
# outputs = self._old_compute_loss(model, inputs, *args, **kwargs)
outputs = self._old_compute_loss(model, inputs, *args, **kwargs)
return outputs
pass