Update superposed/llama/superpose.py

superposed/llama/superpose.py

@@ -198,7 +198,7 @@ class Superpose(nn.Module):
         probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
         next_token = torch.gather(probs_idx, -1, torch.topk(probs_sort, k, dim=-1)[1])       
         # Set all other probs to 0
-        new_probs_map = torch.zeros(probs.shape).bool()
+        new_probs_map = torch.zeros(probs.shape, device="cuda").bool()
         new_probs_map[torch.repeat_interleave(torch.arange(n_prompts), k), torch.flatten(next_token)] = True
         new_probs = torch.where(new_probs_map, probs, 0)
         # Renormalize