custom loss function in Keras combining multiple outputs

I did a lot of searching and am still unable to figure out writing a custom loss function with multiple outputs where they interact.

I have a Neural Network defined as :

def NeuralNetwork():

    inLayer = Input((2,));
    layers = [Dense(numNeuronsPerLayer,activation = 'relu')(inLayer)];
    for i in range(10):
        hiddenLyr = Dense(5,activation = 'tanh',name = "layer"+ str(i+1))(layers[i]);
        layers.append(hiddenLyr);
    out_u = Dense(1,activation = 'linear',name = "out_u")(layers[i]);
    out_k = Dense(1,activation = 'linear',name = "out_k")(layers[i]);

    outLayer = Concatenate(axis=-1)([out_u,out_k]);

    model = Model(inputs = [inLayer], outputs = outLayer);

    return model

I am now trying to define a custom loss function as follows :

def computeLoss(true,prediction):

          u_pred = prediction[:,0];
          k_pred = prediction[:,1];
          loss = f(u_pred)*k_pred;
          return loss;

Where f(u_pred) is some manipulation of u_pred. The code seems to work correct and produce correct results when I use only u_pred (i.e., single output from the neural network only). However, the moment I try to include another output for k_pred and perform the slice of my prediction tensor in the loss function, I start getting wrong results. I feel I am doing something wrong in handling multiple outputs in Keras but am not sure where my mistake lies. Any help on how I may proceed is welcome.

I figured out that you can't just use indexing ( i.e., [:,0] or [:,1] ) to slice tensors in tf. The operation doesn't seem to work. Instead, use the built in function in tensorflow as detailed in https://www.tensorflow.org/api_docs/python/tf/split?version=stable

So the code that worked was:

(u_pred, k_pred) = tf.split(prediction, num_or_size_splits=2, axis=1)