def catch_layer_inputs(model_wrapper, dataloader, device, num_samples=None):
    """Get the layers from model based on model type."""
    num_samples = num_samples or len(dataloader.dataset)
    first_batch = next(iter(dataloader))
    # sequence length
    seqlen = first_batch[0]["input_ids"].shape[1]
    # embedding dimension
    hidden_size = model_wrapper.hidden_size
    # data type
    dtype = next(iter(model_wrapper.model.parameters())).dtype

    # placeholder to save the layer inputs
    inputs = torch.zeros((num_samples, seqlen, hidden_size), dtype=dtype, device=device)
    cache = {"i": 0, "attention_mask": None}
    # additional inputs to the layers
    additional_input = additional_inputs.get(model_wrapper.model_type, [])
    for input_name in additional_input:
        cache[input_name] = None

    # get layers
    layers = model_wrapper.get_layers(False)

    class FirstLayer(torch.nn.Module):
        def __init__(self, module):
            super().__init__()
            self.module = module

        def forward(self, inputs, **kwargs):
            # handle batch dimension
            for batch in range(inputs.shape[0]):
                if cache["i"] >= num_samples:
                    break
                inputs[cache["i"]] = inputs[batch]
                cache["i"] += 1
            cache["attention_mask"] = kwargs.get("attention_mask")
            for input_name in additional_input:
                cache[input_name] = kwargs.get(input_name)
            raise ValueError("Stop forward propagation")

    # put all modules until the first layer on the device
    for module in model_wrapper.get_embeds(False):
        module.to(device)

    # wrap the first layer
    layers[0] = FirstLayer(layers[0])

    # run the model on the data
    for data, _ in dataloader:
        input_ids = data["input_ids"].to(device)
        try:
            model_wrapper.model(input_ids)
        except ValueError:
            pass
        # stop if we have enough samples
        if cache["i"] >= num_samples:
            break

    # unwrap the first layer
    layers[0] = layers[0].module

    # put all modules until the first layer back on the CPU
    for module in model_wrapper.get_embeds(False):
        module.to("cpu")

    if "cuda" in str(device):
        torch.cuda.empty_cache()

    extras = {}
    for input_name in additional_input:
        extras[input_name] = cache[input_name]

    return inputs, cache["attention_mask"], extras
Get the layers from model based on model type.
olive/passes/pytorch/sparsegpt_utils.py