import tensorflow as tf

                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)
            emb_var, other_var = sok.split_embedding_variable_from_others(sok_sparse_demo.trainable_variables)
            grads = tf.gradients(loss, emb_var + other_var, colocate_gradients_with_ops=True,
                                 unconnected_gradients=tf.UnconnectedGradients.NONE)
            emb_grads, other_grads = grads[:len(emb_var)], grads[len(emb_var):]
            if "plugin" in args.optimizer:
                emb_train_op = emb_opt.apply_gradients(zip(emb_grads, emb_var))
            else:
                with sok.OptimizerScope(emb_var):
                    emb_train_op = emb_opt.apply_gradients(zip(emb_grads, emb_var))
            with tf.control_dependencies([*emb_grads]):
                # in case NCCL runs concurrently via SOK and horovod
                other_grads = strategy.reduce("sum", other_grads)
            other_train_op = dense_opt.apply_gradients(zip(other_grads, other_var))

