nnsight
nnsight copied to clipboard
[0.5] Unpacking saved tuple results in non saved variables
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