Update app.py

app.py

@@ -118,7 +118,7 @@ def model_main(args, master_port, rank, request_queue, response_queue, mp_barrie
 
     text_encoder = AutoModel.from_pretrained(
         "google/gemma-2-2b", torch_dtype=dtype, device_map="cuda", token=hf_token
-    ).eval()
+    ).eval().to("cuda")
     cap_feat_dim = text_encoder.config.hidden_size
     if args.num_gpus > 1:
         raise NotImplementedError("Inference with >1 GPUs not yet supported")