diff --git a/test_model_loading.py b/test_model_loading.py index be80840..328f7dc 100644 --- a/test_model_loading.py +++ b/test_model_loading.py @@ -248,6 +248,12 @@ def test_strategy_6(): model, use_gradient_checkpointing=False, ) + print(" ✓ Model prepared for k-bit training") + + # Actually quantize the model + from bitsandbytes.nn.modules import Params4bit + print(" Quantizing weights to 4-bit...") + model.quantize_4bit() print(" ✓ Model quantized to 4-bit (~17.5GB)") print("\n Step 3: Move to GPU 0...")