ValueError when extracting activations

[karnapmk]

I seem to be running into a ValueError when extracting activations with model.trace(). I share code with a collaborator and it seems to run perfectly on his end. For reference, my servers use Tesla V100 GPUs with 32GB of RAM.

from nnsight import NNsight
from transformers import AutoTokenizer, EsmForMaskedLM

model_name = "facebook/esm2_t36_3B_UR50D"
device = "cuda:1"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = EsmForMaskedLM.from_pretrained(model_name, device_map="auto")
model = NNsight(model)
model.tokenizer = tokenizer
model.eval()

sequence = "MDLSALRVEEVQNVINAMQKILECPICLELIKEPVSTKCDHIFCKFCMLKLLNQKKGPSQCPLCKNDITKRSLQESTRFSQLVEELLKIICAFQLDTGLE"
tokens = tokenizer(sequence, return_tensors="pt").to(device)

tracer_kwargs = {'scan': False, 'validate': False}
act_module = model.esm.encoder.layer[18]

with model.trace(**tokens, **tracer_kwargs):
        activations = act_module.output.output.save()

[Butanium]

Can you share your nnsight version and the exact stack trace? Ty!