Add correct device

app.py

@@ -45,6 +45,7 @@ with tab1:
         vectors_2d=vectors_2d, 
         reducer=reducer,
         corrector=corrector,
+        device=device,
     )
 
 with tab2:

views.py

@@ -53,7 +53,7 @@ def diffs(embeddings: np.ndarray, corrector, encoder: PreTrainedModel, tokenizer
 
     # st.html('<a href="https://www.flaticon.com/free-icons/array" title="array icons">Array icons created by Voysla - Flaticon</a>')
 
-def plot(df: pd.DataFrame, embeddings: np.ndarray, vectors_2d, reducer, corrector):
+def plot(df: pd.DataFrame, embeddings: np.ndarray, vectors_2d, reducer, corrector, device):
     
 
     # Add a scatter plot using Plotly
@@ -83,7 +83,7 @@ def plot(df: pd.DataFrame, embeddings: np.ndarray, vectors_2d, reducer, correcto
 
     with col1:
         # Main content stays here (scatterplot, form, etc.)
-        selected_points = plotly_events(fig, click_event=True, hover_event=False, #override_height=600, override_width="100%"
+        selected_points = plotly_events(fig, click_event=True, hover_event=False,# override_height="600", override_width="600"
                                         )
         with st.form(key="form1_main"):
             if selected_points:
@@ -103,7 +103,7 @@ def plot(df: pd.DataFrame, embeddings: np.ndarray, vectors_2d, reducer, correcto
                 inferred_embedding = inferred_embedding.astype("float32")
 
                 inversion_output_text, = vec2text.invert_embeddings(
-                    embeddings=torch.tensor(inferred_embedding).cuda(),
+                    embeddings=torch.tensor(inferred_embedding).to(device),
                     corrector=corrector,
                     num_steps=20,
                 )