import torch


def difference_models_norm_2(model_1, model_2):
    """
    Return the norm 2 difference between the two model parameters. Used in FedProx. 
    """
    tensor_1_backbone = list(model_1["feat_model"].parameters())
    tensor_1_classifier = list(model_1["classifier"].parameters())
    tensor_2_backbone = list(model_2["feat_model"].parameters())
    tensor_2_classifier = list(model_2["classifier"].parameters())
    
    diff_list = [torch.sum((tensor_1_backbone[i] - tensor_2_backbone[i])**2) for i in range(len(tensor_1_backbone))]
    diff_list.extend([torch.sum((tensor_1_classifier[i] - tensor_2_classifier[i])**2) for i in range(len(tensor_1_classifier))])

    norm = sum(diff_list)
    return norm