colpali-document-retrieval-vision-language-models-cloud is buggy

[suhjohn]

[suhjohn]

for pdf in sample_pdfs:
    page_embeddings = []
    dataloader = DataLoader(
        pdf["images"],
        batch_size=2,
        shuffle=False,
        collate_fn=lambda x: process_images(processor, x),
    )
    for batch_doc in tqdm(dataloader):
        batch_doc = {k: v.to(device) for k, v in batch_doc.items()}
        with torch.no_grad():
            with torch.cuda.amp.autocast():
                embeddings_doc = model(**batch_doc)
        page_embeddings.extend(list(torch.unbind(embeddings_doc.cpu())))
    pdf["embeddings"] = page_embeddings

fix

[thomasht86]

Thanks for reporting!
Sorry about that - our tests run without GPU, so we didn't catch that one.
I will add a PR with fix.