def get_sizes(input_size, hidden_size, output_size, layer_idx, nlayers):
    """Get the input and hidden size for a layer"""
    if nlayers == 1:
        return input_size, output_size
    if layer_idx == 0:
        return input_size, hidden_size
    if layer_idx == nlayers - 1:
        return hidden_size, output_size
    return hidden_size, hidden_size