chore: Update TashkeelModelEO and TashkeelModelED loading in app.py

app.py

@@ -24,16 +24,13 @@ def infer_shakkala(input_text):
 tokenizer = TashkeelTokenizer()
 eo_ckpt_path = Path(__file__).parent / 'models/best_eo_mlm_ns_epoch_193.pt'
 
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-print('device:', device)
-
 max_seq_len = 1024
 print('Creating Model...')
 eo_model = TashkeelModelEO(tokenizer, max_seq_len=max_seq_len, n_layers=6, learnable_pos_emb=False)
 ed_model = TashkeelModelED(tokenizer, max_seq_len=max_seq_len, n_layers=3, learnable_pos_emb=False)
 
-eo_model.load_state_dict(torch.load(eo_ckpt_path, map_location=device)).eval().to(device)
-ed_model.load_state_dict(torch.load(eo_ckpt_path, map_location=device)).eval().to(device)
+eo_model.load_state_dict(torch.load(eo_ckpt_path)).eval()
+ed_model.load_state_dict(torch.load(eo_ckpt_path)).eval()
 
 @spaces.GPU()
 def infer_catt(input_text, choose_model):
@@ -41,8 +38,10 @@ def infer_catt(input_text, choose_model):
     batch_size = 16
     verbose = True
     if choose_model == 'Encoder-Only':
+        eo_model.to("cuda")
         output_text = eo_model.do_tashkeel_batch([input_text], batch_size, verbose)
     else:
+        ed_model.to("cuda")
         output_text = ed_model.do_tashkeel_batch([input_text], batch_size, verbose)
 
     return output_text[0]