From 3af9470115fcc5fdc4cc8bcd8315a67e30423ef8 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 30 Dec 2025 14:09:56 +0530 Subject: [PATCH] up --- docs/source/en/training/distributed_inference.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source/en/training/distributed_inference.md b/docs/source/en/training/distributed_inference.md index 22e8a30427b9..b7cd0e20f481 100644 --- a/docs/source/en/training/distributed_inference.md +++ b/docs/source/en/training/distributed_inference.md @@ -263,8 +263,8 @@ def main(): world_size = dist.get_world_size() pipeline = DiffusionPipeline.from_pretrained( - "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, device_map=device - ) + "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16 + ).to(device) pipeline.transformer.set_attention_backend("_native_cudnn") cp_config = ContextParallelConfig(ring_degree=world_size)