import tensorflow as tf

        control_inputs = [restore_op[-1]] if restore_op else None
        with tf.control_dependencies(control_inputs):
            if args.restore_params:
                filepath = r"./embedding_variables"
                op = sok_saver.restore_from_file(embedding_layer.embedding_variable, filepath)
            else:
                op = sok_saver.load_embedding_values(embedding_layer.embedding_variable, init_tensors[i])
            restore_op.append(op)

    loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=True, reduction="none")
    def _replica_loss(labels, logits):
        loss = loss_fn(labels, logits)
        return tf.nn.compute_average_loss(loss, global_batch_size=args.global_batch_size)

    def _train_step(inputs, labels, training):
        def _step_fn(inputs, labels):
            logit, embedding_vector = sok_sparse_demo(inputs, training=training)
            loss = _replica_loss(labels, logit)
