Remove .cpu() instances

app.py

@@ -9,8 +9,8 @@ def top_token_ids(outputs, threshold=-np.inf):
     "Returns the index of the tokens whose score exceeds a threshold, for each output step"
     indexes = []
     for tensor in outputs['scores']:
-        candidates = np.argwhere(tensor.flatten().cpu() > threshold).numpy()[0]
-        ordering_mask = np.argsort(tensor[0][candidates].cpu())
+        candidates = np.argwhere(tensor.flatten() > threshold).numpy()[0]
+        ordering_mask = np.argsort(tensor[0][candidates])
         candidates = candidates[ordering_mask]
         if not isinstance(candidates, np.ndarray):
             indexes.append(np.array([candidates]))
@@ -24,7 +24,7 @@ def plot_word_scores(top_token_ids, outputs, tokenizer, boolq=False, width=600):
         fig.append_trace(
             go.Bar(
                 y=tokenizer.convert_ids_to_tokens(candidates), 
-                x=outputs['scores'][step][0][candidates].cpu(), 
+                x=outputs['scores'][step][0][candidates], 
                 orientation='h'
             ),
             row=step+1, col=1