Update app.py

app.py

@@ -75,7 +75,7 @@ def encode_prompt(prompt_batch, text_encoder, tokenizer, proportion_empty_prompt
             return_tensors="pt",
         )
 
-    
+        print(f"Text Encoder Device: {text_encoder.device}")
         text_input_ids = text_inputs.input_ids
         prompt_masks = text_inputs.attention_mask
 
@@ -125,7 +125,7 @@ def model_main(args, master_port, rank, request_queue, response_queue, mp_barrie
     if args.num_gpus > 1:
         raise NotImplementedError("Inference with >1 GPUs not yet supported")
 
-    tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b", token=hf_token).cuda()
+    tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b", token=hf_token)
     tokenizer.padding_side = "right"
 
     vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", token=hf_token).cuda()