Update app.py

app.py

@@ -5,7 +5,7 @@ import torch
 import os
 from huggingface_hub import login
 from PIL import Image
-from transformers import AutoProcessor, Gemma3ForConditionalGeneration
+from transformers import AutoProcessor, Gemma3ForConditionalGeneration, TextIteratorStreamer
 
 
 """
@@ -18,7 +18,7 @@ duration=None
 login(token = os.getenv('gemma'))
 
 ckpt = "google/gemma-3-4b-it"
-model = Gemma3ForConditionalGeneration.from_pretrained(ckpt, torch_dtype=torch.bfloat16,).to("cuda")
+model = Gemma3ForConditionalGeneration.from_pretrained(ckpt, device_map='cuda', torch_dtype=torch.bfloat16,)
 processor = AutoProcessor.from_pretrained(ckpt)
 
 @spaces.GPU