Update src/flux/pipeline_tools.py

src/flux/pipeline_tools.py

@@ -515,7 +515,7 @@ class CustomFluxPipeline:
         if config["model"].get("dit_quant", "None")=="int8-quanto":
             self.pipe = FluxPipeline.from_pretrained("diffusers/FLUX.1-dev-torchao-int8", 
                                                     torch_dtype=torch_dtype,
-                                                    use_safetensors=False, device_map=device)
+                                                    use_safetensors=False).to(device)
 
         else:
             self.pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev",