# Automatically cast input to model's device map

**URL:** <https://discuss.huggingface.co/t/automatically-cast-input-to-models-device-map/76975>\
**Category:** Beginners\
**Created:** [March 11, 2024, 8:51pm UTC](https://discuss.huggingface.co/t/automatically-cast-input-to-models-device-map/76975 "2024-03-11T20:51:20Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![souryadey](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/souryadey/32/31175_2.png) [@souryadey](https://discuss.huggingface.co/u/souryadey)\
**Post date:** [March 11, 2024, 8:51pm UTC](https://discuss.huggingface.co/t/automatically-cast-input-to-models-device-map/76975/1 "2024-03-11T20:51:20Z")

</div>

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:

```python
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:

```python
# 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?