ValueError: too many values to unpack (expected 2)

I'm trying to extract the attention map from ESM-1b model using the following code:

import torch
import esm

# Check if CUDA is available and set the device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load ESM-1b model
model, alphabet = esm.pretrained.esm1b_t33_650M_UR50S()
model = model.eval().to(device)  # Move the model to GPU

# Prepare data (single sequence or MSA)
sequence = "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"
batch_converter = alphabet.get_batch_converter()
data = batch_converter([sequence])

# Move the data to GPU
data = [item.to(device) for item in data]

# Forward pass through the model
with torch.no_grad():
    results = model(*data, repr_layers=[33], return_contacts=True)

# Extract attention map
attention_map = results["attns"]

# Now you can use the attention map as a feature for your own model

But I'm getting a ValueError: too many values to unpack (expected 2) error. Could you please help me understand what's going wrong?

facebookresearch / esm

ValueError: too many values to unpack (expected 2) #575