fix

app.py

@@ -127,9 +127,12 @@ class RepLlamaModel:
         return model
 
     def encode(self, texts, batch_size=48, **kwargs):
-        self.model = self.model.cuda()
+        # if model is not on cuda, put it there
+        if self.model.device.type != "cuda":
+            self.model = self.model.cuda()
+            
         all_embeddings = []
-        for i in range(0, len(texts), batch_size):
+        for i in tqdm.tqdm(range(0, len(texts), batch_size)):
             batch_texts = texts[i:i+batch_size]
             
             batch_dict = create_batch_dict(self.tokenizer, batch_texts, always_add_eos="last")
@@ -143,7 +146,7 @@ class RepLlamaModel:
                     logger.info(f"Encoded shape: {embeddings.shape}, Norm of first embedding: {torch.norm(embeddings[0]).item()}")
                     all_embeddings.append(embeddings.cpu().numpy())
 
-        self.model = self.model.cpu()
+        # self.model = self.model.cpu()
         return np.concatenate(all_embeddings, axis=0)
 
 def load_corpus_embeddings(dataset_name):