Skip to content

ESMFold2Model.forward() got an unexpected keyword argument 'pocket_feature' when folding receptor+peptide complexes #40

Description

@RzaMsv

I'm getting a TypeError when trying to fold a receptor+peptide complex using the ESMFold2 model loaded via trust_remote_code. The model itself loads fine and tokenization runs without issue, but the actual folding call fails with an unexpected keyword argument error. I ruled out a stale cache by deleting and re-downloading the cached remote code files entirely, and the error persists, so this looks like a real mismatch in the code rather than something on my end. Details below.

Environment: transformers 5.14.1, torch 2.11.0+cu128, CUDA available, single GPU (RTX 6000 Ada). Loaded via AutoModel.from_pretrained("Synthyra/ESMFold2", trust_remote_code=True, dtype=torch.float32). Model loads successfully in about 5 seconds and tokenization runs cleanly, so this does not appear to be a setup or dependency issue on my end.

Steps to reproduce:

import torch
from transformers import AutoModel
model = AutoModel.from_pretrained(
"Synthyra/ESMFold2",
trust_remote_code=True,
dtype=torch.float32,
).eval().cuda()
result = model.fold_protein(
receptor_seq + ":" + peptide_seq,
num_loops=20,
num_sampling_steps=100,
num_diffusion_samples=1,
seed=0,
)

This is the minimal call that triggers the crash, using a standard receptor plus peptide complex through the documented fold_protein API.

Error:

TypeError: ESMFold2Model.forward() got an unexpected keyword argument 'pocket_feature'

This is the exact exception copied from the traceback.

Notes: Folding a single-chain sequence without the colon separator may work fine, though I have not fully confirmed this. The failure appears specific to the two-chain complex path, since I see a UserWarning about no MSA provided for A_0 and A_1 right before the crash, suggesting pocket conditioning is being added specifically for the receptor-peptide complex case. This helps narrow the bug down to the complex-folding code path rather than the whole model.

I also deleted and re-downloaded the entire cached transformers_modules/Synthyra/ESMFold2 directory to rule out a stale or partial mismatch between fastplms_bundle.py and modeling_esmfold2.py. The same error persisted on a fresh pull, so this looks like a genuine mismatch between whatever assembles the forward-call keyword arguments for complex folding and the current ESMFold2Model.forward() signature, rather than a caching issue on my end. This preempts the likely first suggestion to just clear the cache and retry.

Happy to share the exact commit hash of the cached files and a minimal repro fasta if that would help track this down further.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions