Update app.py

app.py

@@ -174,8 +174,7 @@ def run_cad_recode(point_cloud):
     try:
         input_ids = [tokenizer.pad_token_id] * len(point_cloud) + [tokenizer('<|im_start|>')['input_ids'][0]]
         attention_mask = [-1] * len(point_cloud) + [1]
-        if torch.cuda.is_available():
-            model = cad_recode.cuda()
+        model = cad_recode.cuda()
         with torch.no_grad():
             batch_ids = cad_recode.generate(
                 input_ids=torch.tensor(input_ids).unsqueeze(0).to(model.device),
@@ -266,7 +265,8 @@ tokenizer = AutoTokenizer.from_pretrained(
     padding_side='left')
 cad_recode = CADRecode.from_pretrained(
     'filapro/cad-recode',
-    torch_dtype='auto').eval()
+    torch_dtype='auto',
+    attn_implementation='flash_attention_2').eval()
 
 os.environ['TOKENIZERS_PARALLELISM'] = 'False'
 run()