Fix .generate(input_ids=...)

[borzunov]

This PR fixes the following code (a popular way to run .generate()):

import torch
from transformers import AutoTokenizer
from petals import AutoDistributedModelForCausalLM

model_name = "Maykeye/TinyLlama-v0"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False, add_bos_token=False)
model = AutoDistributedModelForCausalLM.from_pretrained(model_name)

inputs = tokenizer("A cat sat on", return_tensors="pt")
outputs = model.generate(**inputs, max_new_tokens=4)
tokenizer.decode(outputs[0])