[docs] no hard-coding cuda (#3270)

* no hard-coding cuda * Update docs/source/usage_guides/big_modeling.md Co-authored-by: Zach Mueller <[email protected]> * update device_type --------- Co-authored-by: Zach Mueller <[email protected]>
huggingface · Dec 11, 2024 · 3e62fbb · 3e62fbb
1 parent cb8b7c6
commit 3e62fbb
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/docs/source/usage_guides/big_modeling.md b/docs/source/usage_guides/big_modeling.md
@@ -21,7 +21,7 @@ This tutorial will show you how to use Big Model Inference in Accelerate and the
 
 ## Accelerate
 
-A typical workflow for loading a PyTorch model is shown below. `ModelClass` is a model that exceeds the GPU memory of your device (mps or cuda).
+A typical workflow for loading a PyTorch model is shown below. `ModelClass` is a model that exceeds the GPU memory of your device (mps or cuda or xpu).
 
 ```py
 import torch
@@ -64,7 +64,8 @@ Now that the model is fully dispatched, you can perform inference.
 
 ```py
 input = torch.randn(2,3)
-input = input.to("cuda")
+device_type = next(iter(model.parameters())).device.type
+input = input.to(device_type)
 output = model(input)
 ```
 
@@ -91,7 +92,8 @@ model = load_checkpoint_and_dispatch(
 )
 
 input = torch.randn(2,3)
-input = input.to("cuda")
+device_type = next(iter(model.parameters())).device.type
+input = input.to(device_type)
 output = model(input)
 ```