Thank you for the replay!
Yes. They work fine separately during the initialization step, and I don’t think I use in-place operation in the similarity_matrix() and custom_loss().
I initialize the models on their own like this:
    print("initializing model_v1")
    for epoch in range(5):

        for i in range(len(splitted_v1)):
            train = splitted_v1[i]
            # ===================forward=====================
            output = model_v1(train.float())
            loss = torch.norm(output-train, p='fro')
            # ===================backward====================
            optimizer_v1.zero_grad()
            loss.backward()
            optimizer_v1.step()
        # ===================log========================
        print('epoch [{}/{}], loss:{:.4f}'
              .format(epoch + 1, 5, loss))
        lr_scheduler_v1.step()

    print("initializing model_v2")
    for epoch in range(5):

        for i in range(len(splitted_v2)):
            train = splitted_v2[i]
            # ===================forward=====================
            output = model_v2(train.float())
            loss = torch.norm(output - train, p='fro')
            # ===================backward====================
            optimizer_v2.zero_grad()
            loss.backward()
            optimizer_v2.step()
        # ===================log========================
        print('epoch [{}/{}], loss:{:.4f}'
              .format(epoch + 1, 5, loss))
        lr_scheduler_v2.step()

Here is similarity_matrix:
def similarity_matrix(y_pred):
    # S = [\hatX_k \dot \hatX_k^T / p_k]
    s = torch.mm(y_pred, y_pred.t())/y_pred.size()[1]
    return s

Here is custom_loss:
def custom_loss(y_pred, y_true, s_v, local_s):
    S = similarity_matrix(y_pred)
    loss = torch.tensor(0, requires_grad=True, dtype=float)
    for i in range(S.size(0)):
        x = S[i] + sys.float_info.epsilon
        temp = x * torch.log(x) + (1 - x) * torch.log(1 - x)
        loss = loss - temp.sum()

    loss3 = loss/y_pred.size(0)
    # frobenius norm || \hatX_k - X_k ||^2_F
    loss1 = torch.norm(y_pred-y_true, p='fro')

    loss2 = torch.norm(s_v-local_s, p='fro')
    total_loss = loss1 + loss2 + loss3
    return total_loss