From 8d6326868d56eda3390d7864f23ce78df65e7ccf Mon Sep 17 00:00:00 2001 From: vik Date: Fri, 29 Dec 2023 20:39:20 -0800 Subject: [PATCH] allow text model to use MPS/GPUs if available --- moondream/text_model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/moondream/text_model.py b/moondream/text_model.py index 31a86c8..a984aa7 100644 --- a/moondream/text_model.py +++ b/moondream/text_model.py @@ -21,7 +21,7 @@ class TextModel: self.model = load_checkpoint_and_dispatch( self.model, f"{model_path}/text_model.pt", - device_map={"": "cpu"}, + device_map="auto", ) self.text_emb = self.model.get_input_embeddings() @@ -48,7 +48,7 @@ class TextModel: assert prompt.count("") == 1 before, after = prompt.split("") embeds.append(self.text_emb(_tokenize(f"{before}"))) - embeds.append(image_embeds) + embeds.append(image_embeds.to(self.model.device)) embeds.append(self.text_emb(_tokenize(f"{after}"))) return torch.cat(embeds, dim=1)