Getting ValueError for building fedAvg from multi-output keras model

I'm trying to federate a keras model which has multiple outputs. There are two separate dense layers that perform a binary classification and a multi-class classification. I am getting the following ValueError when I try to build my federated averaging process tff.learning.build_federated_averaging_process from model_fn(). Following are the code snippets and error information. I am unable to understand what is going wrong and how to resolve it.

ValueError: in user code:

/home/usr/Envs/tf-fed/lib/python3.7/site-packages/tensorflow_federated/python/learning/framework/optimizer_utils.py:387 _compute_local_training_and_client_delta  *
    client_output = client_delta_fn(dataset, initial_model_weights)
/home/usr/Envs/tf-fed/lib/python3.7/site-packages/tensorflow_federated/python/learning/federated_averaging.py:92 reduce_fn  *
    output = model.forward_pass(batch, training=True)
/home/usr/Envs/tf-fed/lib/python3.7/site-packages/tensorflow_federated/python/learning/framework/dataset_reduce.py:28 _dataset_reduce_fn  *
    return dataset.reduce(initial_state=initial_state_fn(), reduce_func=reduce_fn)
/home/usr/Envs/tf-fed/lib/python3.7/site-packages/tensorflow_federated/python/learning/keras_utils.py:365 forward_pass  *
    return self._forward_pass(batch_input, training=training)
/home/usr/Envs/tf-fed/lib/python3.7/site-packages/tensorflow_federated/python/learning/keras_utils.py:357 _forward_pass  *
    metric.update_state(y_true=y_true, y_pred=predictions)
/home/usr/Envs/tf-fed/lib/python3.7/site-packages/tensorflow/python/keras/utils/metrics_utils.py:90 decorated  **
    update_op = update_state_fn(*args, **kwargs)
/home/usr/Envs/tf-fed/lib/python3.7/site-packages/tensorflow/python/keras/metrics.py:176 update_state_fn
    return ag_update_state(*args, **kwargs)
/home/usr/Envs/tf-fed/lib/python3.7/site-packages/tensorflow/python/keras/metrics.py:604 update_state  **
    y_pred = math_ops.cast(y_pred, self._dtype)
/home/usr/Envs/tf-fed/lib/python3.7/site-packages/tensorflow/python/util/dispatch.py:201 wrapper
    return target(*args, **kwargs)
/home/usr/Envs/tf-fed/lib/python3.7/site-packages/tensorflow/python/ops/math_ops.py:920 cast
    x = ops.convert_to_tensor(x, name="x")
/home/usr/Envs/tf-fed/lib/python3.7/site-packages/tensorflow/python/framework/ops.py:1499 convert_to_tensor
    ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
/home/usr/Envs/tf-fed/lib/python3.7/site-packages/tensorflow/python/ops/array_ops.py:1502 _autopacking_conversion_function
    return _autopacking_helper(v, dtype, name or "packed")
/home/usr/Envs/tf-fed/lib/python3.7/site-packages/tensorflow/python/ops/array_ops.py:1438 _autopacking_helper
    return gen_array_ops.pack(elems_as_tensors, name=scope)
/home/usr/Envs/tf-fed/lib/python3.7/site-packages/tensorflow/python/ops/gen_array_ops.py:6477 pack
    "Pack", values=values, axis=axis, name=name)
/home/usr/Envs/tf-fed/lib/python3.7/site-packages/tensorflow/python/framework/op_def_library.py:744 _apply_op_helper
    attrs=attr_protos, op_def=op_def)
/home/usr/Envs/tf-fed/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py:593 _create_op_internal
    compute_device)
/home/usr/Envs/tf-fed/lib/python3.7/site-packages/tensorflow/python/framework/ops.py:3485 _create_op_internal
    op_def=op_def)
/home/usr/Envs/tf-fed/lib/python3.7/site-packages/tensorflow/python/framework/ops.py:1975 __init__
    control_input_ops, op_def)
/home/usr/Envs/tf-fed/lib/python3.7/site-packages/tensorflow/python/framework/ops.py:1815 _create_c_op
    raise ValueError(str(e))

ValueError: Dimension 1 in both shapes must be equal, but are 1 and 3. Shapes are [?,1] and [?,3].
    From merging shape 0 with other shapes. for '{{node Cast_1/x}} = Pack[N=2, T=DT_FLOAT, axis=0](functional_1/eye_output/Sigmoid, functional_1/mouth_output/Softmax)' with input shapes: [?,1], [?,3].

My model_fn() looks like this:

def model_fn():
    losses = [tf.keras.losses.BinaryCrossentropy(), tf.keras.losses.SparseCategoricalCrossentropy()]
    metrics = [tf.keras.metrics.BinaryAccuracy(),tf.keras.metrics.SparseCategoricalAccuracy()]

    keras_model = build_model()
    return tff.learning.from_keras_model(
        keras_model,
        input_spec=spec,
        loss=losses,
        metrics=metrics)

where build_model() creates the keras model:

build_model():
  ...
  out1 = Dense(1, activation='sigmoid')(fc1)
  out2 = Dense(3, activation='softmax')(fc2)
  model =  Model(inputs=inputs, outputs=[out1, out2])
  return model

And input_specification that looks like this

OrderedDict([('x',
          TensorSpec(shape=(None, 240, 320), dtype=tf.float32, name=None)),
         ('y',
          (TensorSpec(shape=(None, 1), dtype=tf.int64, name=None),
           TensorSpec(shape=(None, 1), dtype=tf.int64, name=None)))])

How can I build my TFF fedAvg process using such a model?