More generic datset filter

1 files changed, 2 insertions, 2 deletions

diff --git a/common.py b/common.py
index 8d6b55d..7ffa77f 100644
--- a/common.py
+++ b/common.py
@@ -18,7 +18,7 @@ def load_text_embedding(embeddings, token_id, file):
 
 def load_text_embeddings(tokenizer: CLIPTokenizer, text_encoder: CLIPTextModel, embeddings_dir: Path):
     if not embeddings_dir.exists() or not embeddings_dir.is_dir():
-        return 0
+        return []
 
     files = [file for file in embeddings_dir.iterdir() if file.is_file()]
 
@@ -33,4 +33,4 @@ def load_text_embeddings(tokenizer: CLIPTokenizer, text_encoder: CLIPTextModel,
     for (token_id, file) in zip(token_ids, files):
         load_text_embedding(token_embeds, token_id, file)
 
-    return added
+    return tokens