from transformers import AutoModelForCausalLM, AutoTokenizer
import torch, time

target_id = "facebook/opt-6.7b"
draft_id  = "facebook/opt-125m"

tokenizer = AutoTokenizer.from_pretrained(target_id)

target = AutoModelForCausalLM.from_pretrained(target_id, torch_dtype=torch.float16, device_map="auto")
draft  = AutoModelForCausalLM.from_pretrained(draft_id,  torch_dtype=torch.float16, device_map="auto")