diff --git a/nodes_sampler.py b/nodes_sampler.py index 8112e8f..7e4ff44 100644 --- a/nodes_sampler.py +++ b/nodes_sampler.py @@ -68,6 +68,7 @@ def offload_transformer(transformer): setattr(module, attr_name, MetaParameter(param.data.dtype, quant_type)) else: pass + remove_lora_from_module(transformer) else: transformer.to(offload_device)