Update app.py

app.py

@@ -40,7 +40,7 @@ print_keyword=True
 print_topk_patches = True
 
 torch_dtype = torch.float16
-attn_implementation = 'flash_attention_2'
+attn_implementation = 'sdpa'
 device_map = 'cuda'
 
 conv_template = conv_templates['llama_3']
@@ -56,7 +56,8 @@ config = GeckoConfig.from_pretrained(model,
                                     print_keyword=print_keyword)
 processor = GeckoProcessor.from_pretrained(model, config=config, use_keyword=True, cropping_method=cropping_method, crop_size=crop_size)
 model = GeckoForConditionalGeneration.from_pretrained(
-    model, config=config)
+    model, config=config, torch_dtype=torch_dtype, 
+    attn_implementation=attn_implementation, device_map=device_map)
 model.load_text_encoder(processor)
 
 @spaces.GPU