nnsight icon indicating copy to clipboard operation
nnsight copied to clipboard

[0.5] Unpacking saved tuple results in non saved variables

Open Butanium opened this issue 5 months ago • 0 comments

Would be cool if we could unpack saved shapes into save ints

from nnsight import LanguageModel
model = LanguageModel("gpt2")
with model.trace("hello") as tracer:
    batch_size, seq_len = model.inputs[1]["input_ids"].shape.save()
    # batch_size = batch_size.save()
    # seq_len = seq_len.save()
    print(f"batch_size: {batch_size}, seq_len: {seq_len}")
    
print(f"batch_size: {batch_size}, seq_len: {seq_len}")
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[6], line 7
      4     batch_size, seq_len = model.inputs[1]["input_ids"].shape.save()
      5     print(f"batch_size: {batch_size}, seq_len: {seq_len}")
----> 7 print(f"batch_size: {batch_size}, seq_len: {seq_len}")

NameError: name 'batch_size' is not defined

Butanium avatar Jul 21 '25 16:07 Butanium