Automatically cast input to model's device map
Is there a way to automatically infer the device of the model when using auto device map, and cast the input (prompt IDs) to that?Here’s what I have now:DEVICE = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"

tokenizer = transformers.AutoTokenizer.from_pretrained(<model_id>)
model = transformers.AutoModelForCausalLM.from_pretrained(<model_id>, device_map = 'auto')

prompt = "Hello how are you"
prompt_obj = tokenizer(prompt, return_tensors = 'pt').to(DEVICE)
# proceed with model.generateInstead of hardcodingDEVICE, I’d like to infer it from the model’s device map. Something like:# inferred_device = <some code that maybe involves model.REDACTED_TOKEN_map>
prompt_obj = tokenizer(prompt, return_tensors = 'pt').to(inferred_device)Is there a way to do this?