Automatically cast input to model's device map

Beginners
1

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.generate

Instead of hardcoding DEVICE, I’d like to infer it from the model’s device map. Something like:

# inferred_device = <some code that maybe involves model.hf_device_map>
prompt_obj = tokenizer(prompt, return_tensors = 'pt').to(inferred_device)

Is there a way to do this?