Untitled

model_name = 'meta-llama/Llama-3.2-3B'
model = LanguageModel(model_name, device_map=DEVICE)
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = 'left'

tokenized_inputs = tokenizer.batch_encode_plus(input_prompts, return_tensors="pt", padding=True)["input_ids"]
    
with torch.no_grad():
    with model.generate(tokenized_inputs, max_new_tokens=max_new_tokens, do_sample=False, remote=False) as tracer:
        out = model.generator.output.save()