slice tensor on gpu side

octoml · Jan 30, 2024 · f61e148 · f61e148
1 parent 37579fc
commit f61e148
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/serve/mlc_serve/model/model_common.py b/serve/mlc_serve/model/model_common.py
@@ -155,7 +155,7 @@ def _is_safe_to_sample(prob_like):
         torch.cuda.nvtx.range_pop()
         return None
 
-    res_random = torch.multinomial(probs, 1, True).cpu().numpy()[:, 0]
+    res_random = torch.multinomial(probs, 1, True)[:, 0].cpu().numpy()
 
     if logits_random.shape[0] == num_seq:
         torch.cuda.nvtx.range_pop()