Based on the error message and the line of code (assuming the blocking launch was working in the notebook), I guess you might use an invalid index in:
drop_idxs = torch.topk(A, n_drop, largest=False, sorted=False)[-1]

Run the code on the CPU to get a better error message or, if that doesn’t help, print A and n_drop and check for invalid indices.