Correctly sending to device tensors

app.py

@@ -70,9 +70,9 @@ def text_to_image(image, prompt, duplications: float):
     maurice_embeddings = list()
     for batch in batched_inputs:
         # Load on device
-        batch['input_ids'].to(model.device)
-        batch['attention_mask'].to(model.device)
-        batch['pixel_values'].to(model.device)
+        batch['input_ids'] = batch['input_ids'].to(model.device)
+        batch['attention_mask'] = batch['attention_mask'].to(model.device)
+        batch['pixel_values'] = batch['pixel_values'].to(model.device)
         output = model.generate(**batch, max_new_tokens=500, temperature=0.3)
 
         # Unload GPU
@@ -85,7 +85,7 @@ def text_to_image(image, prompt, duplications: float):
 
         for text in generated_text:
             text_output = text.split("ASSISTANT:")[-1]
-            text_embeddings = embedder.encode(text_output).to('cpu')
+            text_embeddings = embedder.encode(text_output)
             maurice_description.append(text_output)
             maurice_embeddings.append(text_embeddings)