Function weights_to_device

utils/offloading.py:114–119

(layer: nn.Module, device: torch.device)

def weights_to_device(layer: nn.Module, device: torch.device):
    for name, module in layer.named_modules():
        if device.type == 'cpu' and 'lora' in name:
            continue
        if hasattr(module, "weight") and module.weight is not None:
            module.weight.data = module.weight.data.to(device, non_blocking=True)


class Offloader: