Skip to content

Commit 63a5c87

Browse files
sayakpaulstevhliu
andauthored
Update distributed_inference.md to include transformer.device_map (#9553)
* Update distributed_inference.md to include `transformer.device_map` * Update docs/source/en/training/distributed_inference.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> --------- Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
1 parent 1287822 commit 63a5c87

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

docs/source/en/training/distributed_inference.md

+1-1
Original file line numberDiff line numberDiff line change
@@ -177,7 +177,7 @@ transformer = FluxTransformer2DModel.from_pretrained(
177177
```
178178

179179
> [!TIP]
180-
> At any point, you can try `print(pipeline.hf_device_map)` to see how the various models are distributed across devices. This is useful for tracking the device placement of the models.
180+
> At any point, you can try `print(pipeline.hf_device_map)` to see how the various models are distributed across devices. This is useful for tracking the device placement of the models. You can also try `print(transformer.hf_device_map)` to see how the transformer model is sharded across devices.
181181
182182
Add the transformer model to the pipeline for denoising, but set the other model-level components like the text encoders and VAE to `None` because you don't need them yet.
183183

0 commit comments

Comments
 (0)