import torch


def sort_rows(m, n_rows):
    """sort N*M matrix by row
    
    Args:
        m (torch.Tensor): N*M matrix to be sorted
        n_rows (int): no of rows to be sorted
       
    Returns:
        sorted (torch.Tensor): N*M matrix with sorted row
    """
    m_T = m.transpose(1, 0)
    sorted_m = torch.topk(m_T, k=n_rows)[0]
    return sorted_m.transpose(1,0)