Thank you for your response!
I tried running the code on CPU but the problem doesn’t occur there,  So i don’t think the problem is from invalid indexing.
also when I changed the code (for logging purposes) to

def magnitude_pruning(A, drop=0.2):
  if drop<0:
    drop=0
  shape_a = A.shape
  n_elem = A.numel()
  A = A.view(-1)
  n_drop = int(n_elem * drop)
  print("....n_drop = ",n_drop)
  print("....A size = ", n_elem)
  top_k_res = torch.topk(A, n_drop, largest=False, sorted=False)
  print("...topk result = ",top_k_res)

  drop_idxs = top_k_res[1]
  mask = torch.ones(n_elem, device=device)
  mask[drop_idxs] = 0
  A = A * mask 
  A = A.view(shape_a)
  return A, mask.view(shape_a)

The error somewhat changed to this
....n_drop =  3816
....A size =  9408
...topk result =  torch.return_types.topk(
values=tensor([-0.0028, -0.0026, -0.0069,  ..., -0.0078, -0.0095, -0.0018],
       device='cuda:0'),
indices=tensor([   1,    2,    6,  ..., 9401, 9407, 8464], device='cuda:0'))
....n_drop =  13695
....A size =  36864
...topk result =  torch.return_types.topk(
values=tensor([-0.0278, -0.0324, -0.2429,  ..., -0.0029, -0.0083, -0.0025],
       device='cuda:0'),
indices=tensor([    1,     3,     4,  ..., 36851, 36863, 35230], device='cuda:0'))
....n_drop =  0
....A size =  36864
...topk result =  torch.return_types.topk(
values=tensor([], device='cuda:0'),
indices=tensor([], device='cuda:0', dtype=torch.int64))
....n_drop =  5047
....A size =  36864
...topk result =  
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-10-f45d5e9895b6> in <module>()
     10     else:
     11       cur = getattr(cur, attr)
---> 12   W_ = magnitude_pruning(cur.weight.detach(),value)[0]

5 frames
<ipython-input-7-51b1f75557a0> in magnitude_pruning(A, drop)
     10   print("....A size = ", n_elem)
     11   top_k_res = torch.topk(A, n_drop, largest=False, sorted=False)
---> 12   print("...topk result = ",top_k_res)
     13 
     14   drop_idxs = top_k_res[1]

/usr/local/lib/python3.6/dist-packages/torch/tensor.py in __repr__(self)
    177             return handle_torch_function(Tensor.__repr__, relevant_args, self)
    178         # All strings are unicode in Python 3.
--> 179         return torch._tensor_str._str(self)
    180 
    181     def backward(self, gradient=None, retain_graph=None, create_graph=False):

/usr/local/lib/python3.6/dist-packages/torch/_tensor_str.py in _str(self)
    370 def _str(self):
    371     with torch.no_grad():
--> 372         return _str_intern(self)

/usr/local/lib/python3.6/dist-packages/torch/_tensor_str.py in _str_intern(self)
    350                     tensor_str = _tensor_str(self.to_dense(), indent)
    351                 else:
--> 352                     tensor_str = _tensor_str(self, indent)
    353 
    354     if self.layout != torch.strided:

/usr/local/lib/python3.6/dist-packages/torch/_tensor_str.py in _tensor_str(self, indent)
    239         return _tensor_str_with_formatter(self, indent, summarize, real_formatter, imag_formatter)
    240     else:
--> 241         formatter = _Formatter(get_summarized_data(self) if summarize else self)
    242         return _tensor_str_with_formatter(self, indent, summarize, formatter)
    243 

/usr/local/lib/python3.6/dist-packages/torch/_tensor_str.py in __init__(self, tensor)
     87 
     88         else:
---> 89             nonzero_finite_vals = torch.masked_select(tensor_view, torch.isfinite(tensor_view) & tensor_view.ne(0))
     90 
     91             if nonzero_finite_vals.numel() == 0:

RuntimeError: CUDA error: device-side assert triggered

I tried to look into the source of  the problem and found that when the value of drop is very low (~0) is when it occur, so I tried using torch.topk on an arbitrary tensor with k=0, it works for the first time and then the next time it gives this error RuntimeError: CUDA error: device-side assert triggered  , by now I can fix the issue by just adding an if statement, but I’m curious about the reason it happens,
Thank you!