Use correct device

views.py

@@ -22,7 +22,7 @@ def diffs(embeddings: np.ndarray, corrector, encoder: PreTrainedModel, tokenizer
     st.markdown("This application lets you freely explore to which extent that property applies to embedding inversion models given the other factors of inaccuracy")
 
     generated_sentence = ""
-
+    device = encoder.device
 
     with st.form(key="foo") as form:
         submit_button = st.form_submit_button("Synthesize")
@@ -35,10 +35,10 @@ def diffs(embeddings: np.ndarray, corrector, encoder: PreTrainedModel, tokenizer
         st.latex("=")
 
         if submit_button:
-            v1, v2, v3 = get_gtr_embeddings([sent1, sent2, sent3], encoder, tokenizer, device=encoder.device).to("cpu")
+            v1, v2, v3 = get_gtr_embeddings([sent1, sent2, sent3], encoder, tokenizer, device=encoder.device).to(device)
             v4 = v1 - v2 + v3
             generated_sentence, = vec2text.invert_embeddings(
-                embeddings=v4.unsqueeze(0).cuda(),
+                embeddings=v4.unsqueeze(0).to(device),
                 corrector=corrector,
                 num_steps=20,
             )