Update on TopK

Update on TopK code.