diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 745b210208..980425e1f1 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -55,7 +55,7 @@ else: pass # Reduce VRAM usage by reducing fragmentation -os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" +os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True,roundup_power2_divisions:[64:128,256:64,>:32]" # Hugging Face Hub faster downloads if "HF_HUB_ENABLE_HF_TRANSFER" not in os.environ: