From 5d134a3cef78c105525f9ced4e88438496e90656 Mon Sep 17 00:00:00 2001 From: amanturamatov Date: Tue, 6 Jun 2023 14:38:41 +0600 Subject: [PATCH] Solve bug for 'get_text_embedding' with batch_size=1 by not squeezing the tokenizer output --- src/laion_clap/hook.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/laion_clap/hook.py b/src/laion_clap/hook.py index fac512f..a97b6fa 100644 --- a/src/laion_clap/hook.py +++ b/src/laion_clap/hook.py @@ -69,7 +69,7 @@ def tokenizer(self, text): max_length=77, return_tensors="pt", ) - return {k: v.squeeze(0) for k, v in result.items()} + return result def load_ckpt(self, ckpt = None, model_id = -1, verbose = True): """Load the pretrained checkpoint of CLAP model