Custom loss function with multiple outputs in tensorflow

my network has two outputs and single input. I am trying to write a custom loss function

$$Loss = Loss_1(y^{true}_1, y^{pred}_1) + Loss_2(y^{true}_2, y^{pred}_2)$$ I was able to write a custom loss function for a single output. But for multiple output, I am struck. Below I wrote a mwe I tried

def model(input_shape=4, output_shape=3, lr=0.0001):
"""
single input and multi-output
loss = custom_loss(out_1_true, out_1_pred)+mse(out_2_true, out_2_pred))
"""
input_layer = Input(input_shape)
layer_1 = Dense(units=32, activation='elu', kernel_initializer='he_uniform')(input_layer)
#outputs
y_1 = Dense(units=output_shape, activation='softmax', kernel_initializer='he_uniform')(layer_1)
y_2 = Dense(units=1, activation='linear', kernel_initializer='he_uniform')(layer_1)

def custom_loss(y_true, y_pred):

# both true value of out_1, out_2 are encoded in y_true
y_true_1     = y_true[:, :1+output_shape]
y_true_2     = y_true[:, 1+output_shape:]

#(this part is wrong...I dont know how)
y_pred_1, y_pred_2 =  y_pred[:, :1+output_shape], y_pred[:, 1+output_shape:]

#custorm loss for y_pred_1
entropy_loss = -y_pred_1 * K.log(y_true_1 + 1e-10)

#mse for y_pred_2
mse          = -K.mean(K.square(y_pred_2 - y_true_2))

#net loss
loss = entropy_loss + C * mse
return loss

Network_model = Model(inputs = input_layer, outputs = [y_1, y_2])
Network_model.compile(loss = custom_loss, optimizer=RMSprop(lr=lr))

return Network_model


I think, the main issue lies in spiting the y_pred tensor.

PS: For the purpose of mwe, I have use normal cross entropy loss and mse loss functions in the above code. But however, I have a different cost functions.

def model(input_shape=4, output_shape=3, lr=0.0001):
"""
single input and multi-output
loss = custom_loss(out_1_true, out_1_pred)+mse(out_2_true, out_2_pred))
"""
input_layer = Input(input_shape)
layer_1 = Dense(units=32, activation='elu', kernel_initializer='he_uniform')(input_layer)

y_1 = Dense(units=output_shape, activation='softmax', kernel_initializer='he_uniform', name='actions')(layer_1)
y_2 = Dense(units=1, activation='linear', kernel_initializer='he_uniform', name='values')(layer_1)

# STEP-2:    Define the custom custom_loss. Note that the args should be y_true, y_pred
def custom_loss(y_true, y_pred):
entropy_loss = -y_pred * K.log(y_true + 1e-10)
return entropy_loss

# STEP-3:   Define the dicts for loss and lossweights with keys as the name of the output layers
LossFunc    =     {'actions':custom_loss, 'values':'mse'}
lossWeights =     {'actions':0.5, 'values':0.5}

Network_model = Model(inputs = input_layer, outputs = [y_1, y_2])

# STEP-4:   complie using LossFunc and lossWeights dicts
Network_model.compile(optimizer=RMSprop(lr=lr), loss=LossFunc, loss_weights=lossWeights, metrics=["accuracy"])

return Network_model

#training example:

#model
Network_model = model(4,2)
# input
S = np.reshape([1.,2.,3.,4.], (1, -1))
#true outputs
action, value = np.reshape([0.1], (1, -1)), np.reshape([0.1, 0.9] , (1, -1))
Y = [action, value]

Network_model.fit(S, [A_t, V_t])

1/1 [==============================] - 0s 2ms/step - loss: 6.7340 - actions_loss: 1.1512 - values_loss: 12.3169 - actions_accuracy: 0.0000e+00 - values_accuracy: 1.0000