Try something

Author: larryliu0820

quantization/quantize.py

@@ -63,6 +63,7 @@ def quantize_model(model: nn.Module, device, quantize_options, tokenizer=None):
             # Only use quant API for dtype bf16 and CUDA
             if quantizer == "linear:int4" and precision == torch.bfloat16 and device == "cuda":
                 quantize_(model, int4_weight_only(group_size=q_kwargs["groupsize"]))
+                model.to(device=device)
                 continue
 
             try: