Thanks @BramVanroy for helping out with this. I guess I figured out the batching part once I went through the documentation a bit more carefully. I do however have a followup:
I moved the models and inputs to GPU hoping that would be faster, but it seems that doesn’t really work with multiprocessing?
Could you tell me if there’s something really obvious I’m missing here?
from transformers import DPRContextEncoder, DPRContextEncoderTokenizer

import torch

torch.set_grad_enabled(False)

ctx_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")

ctx_tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")

ctx_encoder = ctx_encoder.to('cuda:0')

from datasets import load_from_disk
# a version of the MARCO dataset that only has passage text
ds = load_from_disk("drive/MyDrive/marco")
print(ctx_encoder.device)

def encode(example):
    npys = ctx_encoder(**ctx_tokenizer(example['passages'], return_tensors="pt", padding="longest", truncation="longest_first").to('cuda:0'))[0].cpu().numpy()
return {'embeddings': npys}

ds_with_embeddings = ds.map(encode, batched=True, batch_size=100, num_proc=6)

This fails with the following error on a GPU colab device:
RuntimeError: CUDA error: initialization error
Any help would be appreciated! Thanks a lot!