Thank you so much for open-sourcing the great ESMFold2 model. I was really excited to try it out via HuggingFace!

I am currently attempting to run a toy example (predicting a small number of proteins) using a CPU backend to check the implementation. However, I ran into a RuntimeError due to a data type (dtype) mismatch.

I tried to resolve this by keeping the model in float32 or casting it entirely to torch.bfloat16 via .to(torch.bfloat16), but neither worked. It seems like different parts of the pipeline have conflicting tensor initializations on CPU.

import torch
from transformers.models.esmfold2.modeling_esmfold2 import ESMFold2Model

# Ubiquitin (PDB 1UBQ)
sequence = "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG"

# Load model to CPU
model = ESMFold2Model.from_pretrained("biohub/ESMFold2-Fast").cpu().eval()

# Case 1: Running with default float32 model
output = model.infer_protein(sequence, num_loops=3, num_sampling_steps=50)

Case 1 Error Log (Default float32)

When running the default model, the internal language model (self.language_model) outputs bfloat16 tensors, causing a clash with downstream float32 linear layers:

File ".../transformers/models/esmfold2/modeling_esmfold2_common.py", line 2170, in forward
    lm_z = self.base_z_linear(hidden_states)  # [B, L, 81, d_z]
File ".../torch/nn/modules/linear.py", line 125, in forward
    return F.linear(input, self.weight, self.bias)
RuntimeError: expected m1 and m2 to have the same dtype, but got: c10::BFloat16 != float

Case 2 Error Log (After .to(torch.bfloat16))

When explicitly casting the model using model.to(torch.bfloat16), some internal input embeddings/features (atom_feats) are still generated as standard float32, causing the inverse error:

model.to(torch.bfloat16)
output = model.infer_protein(sequence, num_loops=3, num_sampling_steps=50)

File ".../transformers/models/esmfold2/modeling_esmfold2_common.py", line 863, in forward
    c_base = self.atom_norm(self.atom_linear(atom_feats))
File ".../torch/nn/modules/linear.py", line 125, in forward
    return F.linear(input, self.weight, self.bias)
RuntimeError: expected m1 and m2 to have the same dtype, but got: float != c10::BFloat16

Additional Note (LU Decomposition Error)

As a workaround, when I wrapped the inference with torch.amp.autocast(device_type="cpu", dtype=torch.bfloat16), it bypassed the linear layer errors but then threw a NotImplementedError: "lu_cpu" not implemented for 'BFloat16' inside torch.linalg.det during the structure head sampling phase.

Any advice or official fixes on how to run this smoothly on a CPU environment would be greatly appreciated!

Thank you so much for your time and help!