import torch


def transfer_weights(tf_model, torch_model):
    """ Example function on how to transfer weights from tensorflow model to pytorch model
    
    Parameters
    ----------
    tf_model : tensorflow model
        Weight donor
    torch_model : pytorch model (torch.nn.Modules)
        Weight recipient
    
    Returns
    -------
    torch_model : pytorch model
        Model with updated weights
    
    """
    def convert_weight(W):
        return torch.tensor(W.transpose(3, 2, 0, 1), dtype=torch.float)

    def convert_bias(B):
        return torch.tensor(B, dtype=torch.float)

    all_tf_weights = tf_model.get_weights()
    
    torch_model.conv1.conv.weight.data = convert_weight(all_tf_weights[0])
    torch_model.conv1.conv.bias.data = convert_bias(all_tf_weights[1])

    torch_model.conv2.conv.weight.data = convert_weight(all_tf_weights[2])
    torch_model.conv2.conv.bias.data = convert_bias(all_tf_weights[3])

    torch_model.conv3.conv.weight.data = convert_weight(all_tf_weights[4])
    torch_model.conv3.conv.bias.data = convert_bias(all_tf_weights[5])

    torch_model.conv4.conv.weight.data = convert_weight(all_tf_weights[6])
    torch_model.conv4.conv.bias.data = convert_bias(all_tf_weights[7])

    torch_model.conv5.conv.weight.data = convert_weight(all_tf_weights[8])
    torch_model.conv5.conv.bias.data = convert_bias(all_tf_weights[9])
    
    return torch_model