diff --git a/src/llmcompressor/pipelines/sequential/helpers.py b/src/llmcompressor/pipelines/sequential/helpers.py index e13144689..273ed76d4 100644 --- a/src/llmcompressor/pipelines/sequential/helpers.py +++ b/src/llmcompressor/pipelines/sequential/helpers.py @@ -538,6 +538,8 @@ def dispatch_for_sequential(model: PreTrainedModel) -> PreTrainedModel: if torch.cuda.is_available(): offloaded_dispatch(model, execution_device=torch.device("cuda:0")) + elif hasattr(torch, "npu") and torch.npu.is_available(): + offloaded_dispatch(model, execution_device=torch.device("npu:0")) elif hasattr(torch, "xpu") and torch.xpu.is_available(): offloaded_dispatch(model, execution_device=torch.device("xpu:0")) else: