def router_loss_func_test(gate_logits: torch.Tensor, top_k=2) -> float:
    if gate_logits is None:
        return 0

    if isinstance(gate_logits, tuple):
        # cat along the layers?
        gate_logits = torch.cat(gate_logits, dim=0) # batch_size * num_hidden_layers, sequence_length, num_experts

    num_experts = gate_logits.shape[-1]

    _, expert_indicies = torch.topk(gate_logits, top_k, dim=-1)  # this is done so you don't need to pass expert_indicies
    routing_probs = gate_logits.softmax(dim=-1) # routing probs

    if expert_indicies.dtype != torch.int64: # cast the expert indices to int64, otherwise one-hot encoding will fail
        expert_indicies = expert_indicies.to(torch.int64)

    if len(expert_indicies.shape) == 2:
        expert_indicies = expert_indicies.unsqueeze(2)

    expert_mask = torch.nn.functional.one_hot(expert_indicies, num_experts)

    # For a given token, determine if it was routed to a given expert.
    expert_mask = torch.max(expert_mask, axis=-2).values

    expert_mask = expert_mask.to(torch.float32) # cast to float32 otherwise mean will fail
    tokens_per_group_and_expert = torch.mean(expert_mask, axis=-2)

    router_prob_per_group_and_expert = torch.mean(routing_probs, axis=-2)

    loss = torch.mean(tokens_per_group_and_expert * router_prob_per_group_and_expert) * (num_experts**2)
    return loss