Update README.md

README.md

@@ -32,7 +32,7 @@ def load_quanto_transformer(repo_path):
     with torch.device("meta"):
         transformer = diffusers.FluxTransformer2DModel.from_config(hf_hub_download(repo_path, "transformer/config.json")).to(torch.bfloat16)
     state_dict = load_file(hf_hub_download(repo_path, "transformer/diffusion_pytorch_model.safetensors"))
-    requantize(transformer, state_dict, quantization_map, device=torch.device("cpu"))
+    requantize(transformer, state_dict, quantization_map, device=torch.device("cuda"))
     return transformer
 
 
@@ -44,7 +44,7 @@ def load_quanto_text_encoder_2(repo_path):
     with torch.device("meta"):
         text_encoder_2 = transformers.T5EncoderModel(t5_config).to(torch.bfloat16)
     state_dict = load_file(hf_hub_download(repo_path, "text_encoder_2/model.safetensors"))
-    requantize(text_encoder_2, state_dict, quantization_map, device=torch.device("cpu"))
+    requantize(text_encoder_2, state_dict, quantization_map, device=torch.device("cuda"))
     return text_encoder_2