Function calibrate_model

examples/mediatek/model_export_scripts/whisper.py:376–393

(model, cal_dataset, chunk_idx: str)

def calibrate_model(model, cal_dataset, chunk_idx: str):
    with torch.no_grad():
        for inp in tqdm(cal_dataset, desc="Calibrating Model: "):
            # pass prompt and response
            for batch in tqdm(inp[chunk_idx].keys(), desc="Batch: "):
                if inp[chunk_idx][batch] is not None:
                    inputs_embeds = torch.tensor(inp[chunk_idx][batch]["hidden_state"])
                    mask = torch.tensor(inp[chunk_idx][batch]["mask"])
                    pos_emb = torch.tensor(inp[chunk_idx][batch]["pos_emb"])
                    cache = torch.tensor(inp[chunk_idx][batch]["cache"])
                    cross_cache = torch.tensor(inp[chunk_idx][batch]["cross_cache"])
                    model(
                        inputs_embeds,
                        mask,
                        pos_emb,
                        cross_cache,
                        *torch.split(cache, 1, dim=0),
                    )


def export_to_et_ir(