Update app.py

app.py

@@ -117,7 +117,7 @@ def model_main(args, master_port, rank, request_queue, response_queue, mp_barrie
     dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[args.precision]
 
     text_encoder = AutoModel.from_pretrained(
-        "google/gemma-2-2b", torch_dtype=dtype, device_map="cuda", token=hf_token
+        "google/gemma-2-2b", torch_dtype=dtype, token=hf_token
     ).eval().to("cuda")
     cap_feat_dim = text_encoder.config.hidden_size
     if args.num_gpus > 1: