def permute_to_N_HWA_K(tensor, K: int):
    """
    Transpose/reshape a tensor from (N, (Ai x K), H, W) to (N, (HxWxAi), K)
    """
    assert tensor.dim() == 4, tensor.shape
    N, _, H, W = tensor.shape
    tensor = tensor.view(N, -1, K, H, W)
    tensor = tensor.permute(0, 3, 4, 1, 2)
    tensor = tensor.reshape(N, -1, K)  # Size=(N,HWA,K)
    return tensor

    Transpose/reshape a tensor from (N, (Ai x K), H, W) to (N, (HxWxAi), K)
    
iGPT/models/grit_src/third_party/CenterNet2/detectron2/modeling/meta_arch/dense_detector.py