import tensorflow as tf

    with graph.as_default():
        tf_sparse_demo = TFDemo(vocabulary_size=args.max_vocabulary_size_per_gpu * args.gpu_num,
                                embedding_vec_size=args.embedding_vec_size,
                                combiner=args.combiner,
                                slot_num=args.slot_num,
                                max_nnz=args.max_nnz,
                                use_hashtable=args.use_hashtable,
                                num_of_dense_layers=0)
        
        optimizer = utils.get_dense_optimizer(args.optimizer)(learning_rate=0.1)

        loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=True)
        def _train_step(inputs, labels, training):
            logit, embedding_vector = tf_sparse_demo(inputs, training=training)
            loss = loss_fn(labels, logit)
            grads = tf.gradients(loss, tf_sparse_demo.trainable_variables,
                                 colocate_gradients_with_ops=True,
                                 unconnected_gradients=tf.UnconnectedGradients.NONE)
            train_op = optimizer.apply_gradients(zip(grads, tf_sparse_demo.trainable_variables))
            with tf.control_dependencies([train_op]):
                loss = tf.identity(loss)
                return loss, embedding_vector


        dataset = utils.tf_dataset(*random_samples, batchsize=args.global_batch_size,
                                   to_sparse_tensor=True, repeat=1)
        train_iterator = dataset.make_initializable_iterator()
        iterator_init = train_iterator.initializer

        inputs, labels = train_iterator.get_next()
