diff --git a/README.md b/README.md
index 8abaf31e2..0f9590249 100644
--- a/README.md
+++ b/README.md
@@ -5,8 +5,8 @@
 [![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](https://opensource.org/licenses/MIT)
 [![PyPI Status Badge](https://badge.fury.io/py/elasticdl-client.svg)](https://pypi.org/project/elasticdl-client/)
 
-ElasticDL is a Kubernetes-native deep learning framework built on top of
-TensorFlow 2.0 that supports fault-tolerance and elastic scheduling.
+ElasticDL is a Kubernetes-native deep learning framework
+that supports fault-tolerance and elastic scheduling.
 
 ## Main Features
 
@@ -16,11 +16,11 @@ Through Kubernetes-native design, ElasticDL enables fault-tolerance and works
 with the priority-based preemption of Kubernetes to achieve elastic scheduling
 for deep learning tasks.
 
-### TensorFlow 2.0 Eager Execution
+### Support TensorFlow and PyTorch
 
-A distributed deep learning framework needs to know local gradients before the
-model update. Eager Execution allows ElasticDL to do it without hacking into the
-graph execution process.
+- TensorFlow Estimator.
+- TensorFlow Keras.
+- PyTorch
 
 ### Minimalism Interface
 
@@ -37,30 +37,27 @@ elasticdl train \
   --volume="host_path=/data,mount_path=/data"
 ```
 
-### Integration with SQLFlow
-
-ElasticDL will be integrated seamlessly with SQLFlow to connect SQL to
-distributed deep learning tasks with ElasticDL.
-
-```sql
-SELECT * FROM employee LABEL income INTO my_elasticdl_model
-```
-
 ## Quick Start
 
 Please check out our [step-by-step tutorial](docs/tutorials/get_started.md) for
 running ElasticDL on local laptop, on-prem cluster, or on public cloud such as
 Google Kubernetes Engine.
 
+[TensorFlow Estimator on MiniKube](docs/tutorials/elasticdl_estimator.md)
+
+[TensorFlow Keras on MiniKube](docs/tutorials/elasticdl_local.md)
+
+[PyTorch on MiniKube](docs/tutorials/elasticdl_torch.md )
+
 ## Background
 
-TensorFlow has its native distributed computing feature that is
+TensorFlow/PyTorch has its native distributed computing feature that is
 fault-recoverable. In the case that some processes fail, the distributed
 computing job would fail; however, we can restart the job and recover its status
 from the most recent checkpoint files.
 
-ElasticDL, as an enhancement of TensorFlow's distributed training feature,
-supports fault-tolerance. In the case that some processes fail, the job would
+ElasticDL supports fault-tolerance during distributed training.
+In the case that some processes fail, the job would
 go on running. Therefore, ElasticDL doesn't need to save checkpoint nor recover
 from checkpoints.
 
@@ -80,11 +77,11 @@ first job completes.  In this case, the overall utilization is 100%.
 
 The feature of elastic scheduling of ElasticDL comes from its Kubernetes-native
 design -- it doesn't rely on Kubernetes extensions like Kubeflow to run
-TensorFlow programs; instead, the master process of an ElasticDL job calls
+TensorFlow/PyTorch programs; instead, the master process of an ElasticDL job calls
 Kubernetes API to start workers and parameter servers; it also watches events
 like process/pod killing and reacts to such events to realize fault-tolerance.
 
-In short, ElasticDL enhances TensorFlow with fault-tolerance and elastic
+In short, ElasticDL enhances TensorFlow/PyTorch with fault-tolerance and elastic
 scheduling in the case that you have a Kubernetes cluster. We provide a tutorial
 showing how to set up a Kubernetes cluster on Google Cloud and run ElasticDL
 jobs there.  We respect TensorFlow's native distributed computing feature, which
diff --git a/docs/designs/dynamic_sharding_for_images.md b/docs/designs/dynamic_sharding_for_images.md
new file mode 100644
index 000000000..2cc00b408
--- /dev/null
+++ b/docs/designs/dynamic_sharding_for_images.md
@@ -0,0 +1,271 @@
+# Dynamic Sharding Support Reading Original Images
+
+Now, users have to convert images into RecordIO for ElasticDL. In practice,
+many users read the original images to train their models. So, we need to
+support reading the original data from the storage. However, there are various
+file formats for the original data. We need to design a common definition
+for different formats.
+
+## Different Ways to Store Images and Annotations
+
+1. All images are in the same folder.
+
+    ```txt
+    |-- images
+        |--0001.png
+        |--0002.png
+        |--0003.png
+        |--0004.png
+        |--0005.png
+    ```
+
+    In the format, users may not need labels in the image compression.
+
+2. Images with the same label are in the same folder such as ImageNet.
+
+    ```txt
+    |-- images
+        |--0
+            |--0001.png
+            |--0002.png
+            |--0003.png
+        |--1
+            |--0004.png
+            |--0005.png
+            |--0006.png
+    ```
+
+    Besides the images, there is usually a file to store all filenames and labels
+    like:
+
+    ```csv
+    0001.png,0
+    0002.png,0
+    0003.png,0
+    0004.png,1
+    0005.png,1
+    0006.png,1
+    ```
+
+    Users will read the content from the file to read images from the storage.
+
+3. The description of images, labels, and annotations is in a JSON or XML file.
+
+    For example, the description of the image is in a JSON file for COCO
+    dataset and in a XML file for Pascal VOC dataset. The example of COCO
+    description is
+
+    ```json
+    "{'license': 3,
+    'file_name': 'COCO_val2014_000000391895.jpg',
+    'coco_url': 'http://images.cocodataset.org/val2014/COCO_val2014_000000391895.jpg',
+    'height': 360,
+    'width': 640,
+    'date_captured': '2013-11-14 11:18:45',
+    'flickr_url': 'http://farm9.staticflickr.com/8186/8119368305_4e622c8349_z.jpg',
+    'id': 391895}"
+
+    ```
+
+    ```json
+    "{'image_id': 203564,
+    'id': 37,
+    'caption': 'A bicycle replica with a clock as the front wheel.'}"
+    ```
+
+    In elastic training, we need to split the training data into shards and
+    assign those shards to workers. When the worker fails, we need to
+    reassign uncompleted shards of the worker to other alive workers.
+
+## Shard Definition
+
+In ElasticDL, we define the shard information `start`, `end` and `shard_name`
+in the task. We can define the shard independently and and expose the
+shard for users to create their datasets.
+
+```proto
+message shard {
+  // The storage name for data, such as a MaxCompute Table,
+  // CSV file with image paths.
+  string name = 1;
+
+  // Starting and ending (non-inclusive) record number.
+  int64 start = 2;
+  int64 end = 3;
+}
+```
+
+In order to split the training data into shards, we must get the size of the
+training data firstly. We can get the size by
+`len(os.listdir("images"))` for case 1 and reading the CSV file for case 2.
+For simplicity, we also
+can store the image names in a CSV file for the first way to store images.
+So, the name of the shard is the CSV file. The start and end indices are
+the line number of the CSV file. When the worker gets the shard, it can read
+the images by the lines in the CSV file.
+
+For case 3, it is difficult to get the size  because we don't know the
+format of the description. So, users need to indicate the size of
+training data. We can design a function in the model definition and
+users implement the function to return the size.
+
+```python
+from pycocotools.coco import COCO
+
+def get_training_size():
+    coco = COCO("annotations/captions_val2014.json")
+    return len(coco.anns.keys())
+```
+
+Then, users should define the `training_data` as the Python file with the
+function.
+
+```bash
+--training_data="/test_data/coco/train/create_shard.py"
+```
+
+The `PythonCustomReader` in ElasticDL can load the function to
+get the total size in the master.
+
+```python
+class PythonCustomReader(object):
+     def __init__(self, records_per_task):
+        """
+        Args:
+            kwargs should contains "records_per_task".
+        """
+        AbstractDataReader.__init__(self, **kwargs)
+        self._filename = self._kwargs["filename"]
+        self._records_per_task = self._kwargs["records_per_task"]
+        self._get_size_fn = None
+
+    def load_get_size_fn(self, fn_name="get_training_size"):
+        module = load_module(self._filename)
+        self._get_size_fn = module[fn_name]
+
+    def get_size(self):
+        if self._get_size_fn:
+            return self._get_size_fn()
+
+    def create_shards(self):
+        shard_name_prefix = "shard_"
+        size = self.get_size()
+        shards = {}
+        num_shards = size // self._records_per_task
+        start_ind = 0
+        for shard_id in range(num_shards):
+            shards[shard_name_prefix + str(shard_id)] = (
+                start_ind,
+                self._records_per_task,
+            )
+            start_ind += self._records_per_task
+        return shards
+```
+
+Then, the master will call the function to get the size and split the
+training data into shards by the size. The shard message will only contains
+the start and end index. Users need to read the image information according
+to the index by themselves.
+
+## The Worker Creates the Dataset using Shards
+
+### APIs to Fetch Shards
+
+```python
+class DataShardService(object):
+    def __init__(self, batch_size, master_client=None,):
+        self._mc = master_client
+        if not self._mc
+            master_addr = os.getenv("MASTER_ADDR")
+            worker_id = os.getenv("WORKER_ID")
+            self._mc = MasterClient(
+                build_channel(master_addr), worker_id
+            )
+        self._pending_tasks = []
+        self.record_count = 0
+
+    def fetch_shard(self):
+        return shard
+
+    def report_batch_done(self):
+        if task_done:
+            report_task()
+```
+
+### Create Dataset Using TensorFlow
+
+```python
+import tensorflow as tf
+
+global data_shard_service = DataShardService()
+
+class DynamicShardingHook(tf.train.SessionRunHook):
+    def __init__(self, num_worker):
+        self._max_steps = max_steps
+        self._local_step = 0
+        self._batch_size = 256
+
+    def after_run(self, run_context, run_values):
+        self._local_step += 1
+        if self._local_step * self._batch_size > data_shard_service.record_count:
+            data_shard_service.report_batch_done()
+
+def get_dataset(shuffle=False):
+    def _record_generator():
+        while True:
+            shard = data_shard_service.fetch_shard()
+            if not shard:
+                break
+            records = read_records(shard.start, shard.end)
+            if shuffle:
+                np.random.shuffle(records)
+            for record in records:
+                yield record
+    return tf.data.Dataset.from_generator(_record_generator()
+```
+
+### Create Dataset Using PyTorch
+
+Here, we create the dataset using COCO dataset.
+
+```python
+import torch
+import cv2
+from pycocotools.coco import COCO
+
+global data_shard_service = DataShardService()
+
+coco = COCO("annotations/captions_val2014.json")
+ids = list(coco.anns.keys())
+
+def read_images(shard):
+    images = []
+    for index in range(shard.start, shard.end):
+        ann_id = ids[index]
+        caption = coco.anns[ann_id]['segmentation']
+        img_id = coco.anns[ann_id]['image_id']
+        path = coco.loadImgs(img_id)[0]['file_name']
+        image = cv2.imread(image_path)
+        images.append(image, caption)
+    return images
+
+
+class ImageDataset(torch.utils.data.IterableDataset):
+
+    def __init__(self, shuffle=False):
+        self._shuffle = shuffle
+
+    def __iter__(self):
+        while True:
+            shard = data_shard_service.fetch_shard()
+            if shard:
+                images = read_images(shard)
+                if self._shuffle:
+                    np.random.shuffle(images)
+                for image in images:
+                    yield image
+            else:
+                break
+
+data_loader = DataLoader(dataset=dataset, batch_size=32)
+```
diff --git a/docs/tutorials/elasticdl_deepctr_estimator.md b/docs/tutorials/elasticdl_deepctr_estimator.md
new file mode 100644
index 000000000..d9e1fb383
--- /dev/null
+++ b/docs/tutorials/elasticdl_deepctr_estimator.md
@@ -0,0 +1,173 @@
+# Distributed Training of DeepCTR Estimator using ElasticDL on Kubernetes
+
+This document shows how to run a distributed training job of a deepctr
+estimator model (DeepFM) using [ElasticDL](https://github.com/sql-machine-learning/elasticdl)
+on Kubernetes.
+
+## Prerequisites
+
+1. Install Minikube, preferably >= v1.11.0, following the installation
+   [guide](https://kubernetes.io/docs/tasks/tools/install-minikube).  Minikube
+   runs a single-node Kubernetes cluster in a virtual machine on your personal
+   computer.
+
+1. Install Docker CE, preferably >= 18.x, following the
+   [guide](https://docs.docker.com/docker-for-mac/install/) for building Docker
+   images containing user-defined models and the ElasticDL framework.
+
+1. Install Python, preferably >= 3.6, because the ElasticDL command-line tool is
+   in Python.
+
+## Models
+
+In this tutorial, we use a [DeepFM estimator](https://github.com/shenweichen/DeepCTR/blob/master/deepctr/estimator/models/deepfm.py)
+model in DeepCTR. The complete program to train the model with the
+dataset definition is in [ElasticDL model zoo](https://github.com/sql-machine-learning/elasticdl/tree/develop/model_zoo/deepctr).
+
+## Dataset
+
+In this tutorial, We use the [criteo dataset](https://github.com/shenweichen/DeepCTR/blob/master/examples/criteo_sample.txt)
+in DeepCTR examples.
+
+```bash
+mkdir ./data
+wget https://github.com/shenweichen/DeepCTR/blob/master/examples/criteo_sample.txt -O ./data/criteo_sample.txt
+```
+
+## The Kubernetes Cluster
+
+The following command starts a Kubernetes cluster locally using Minikube.  It
+uses [VirtualBox](https://www.virtualbox.org/), a hypervisor coming with
+macOS, to create the virtual machine cluster.
+
+```bash
+minikube start --vm-driver=virtualbox \
+  --cpus 2 --memory 6144 --disk-size=50gb 
+eval $(minikube docker-env)
+```
+
+The command `minikube docker-env` returns a set of Bash environment variable
+to configure your local environment to re-use the Docker daemon inside
+the Minikube instance.
+
+The following command is necessary to enable
+[RBAC](https://kubernetes.io/docs/reference/access-authn-authz/rbac/) of
+Kubernetes.
+
+```bash
+kubectl apply -f \
+  https://raw.githubusercontent.com/sql-machine-learning/elasticdl/develop/elasticdl/manifests/elasticdl-rbac.yaml
+```
+
+If you happen to live in a region where `raw.githubusercontent.com` is banned,
+you might want to Git clone the above repository to get the YAML file.
+
+## Install ElasticDL Client Tool
+
+The following command installs the command line tool `elasticdl`, which talks to
+the Kubernetes cluster and operates ElasticDL jobs.
+
+```bash
+pip install elasticdl_client
+```
+
+## Build the Docker Image with Model Definition
+
+Kubernetes runs Docker containers, so we need to put user-defined models,
+the ElasticDL api package and all dependencies into a Docker image.
+
+In this tutorial, we use a complete program using a DeepFM estimator model of DeepCTR
+in the ElasticDL repository. To retrieve the source code, please run the following command.
+
+```bash
+git clone https://github.com/sql-machine-learning/elasticdl
+```
+
+Complete codes are in directory [elasticdl/model_zoo/deepctr](https://github.com/sql-machine-learning/elasticdl/tree/develop/model_zoo/deepctr).
+
+We build the image based on tensorflow:1.13.2 and the dockerfile
+is
+
+```text
+FROM tensorflow/tensorflow:1.13.2-py3 as base
+
+RUN pip install elasticdl_api
+RUN pip install deepctr
+
+COPY ./model_zoo model_zoo
+```
+
+Then, we use docker to build the image
+
+```bash
+docker build -t elasticdl:deepctr_estimator -f ${deepctr_dockerfile} .
+```
+
+## Submit the Training Job
+
+The following command submits a training job:
+
+```bash
+elasticdl train \
+  --image_name=elasticdl/elasticdl:1.0.0 \
+  --worker_image=elasticdl:deepctr_estimator \
+  --ps_image=elasticdl:deepctr_estimator \
+  --job_command="python -m model_zoo.deepctr.deepfm_estimator --training_data=/data/criteo_sample.txt --validation_data=/data/criteo_sample.txt" \
+  --num_workers=1 \
+  --num_ps=1 \
+  --num_evaluator=1 \
+  --master_resource_request="cpu=0.2,memory=1024Mi" \
+  --master_resource_limit="cpu=1,memory=2048Mi" \
+  --ps_resource_request="cpu=0.2,memory=1024Mi" \
+  --ps_resource_limit="cpu=1,memory=2048Mi" \
+  --worker_resource_request="cpu=0.3,memory=1024Mi" \
+  --worker_resource_limit="cpu=1,memory=2048Mi" \
+  --chief_resource_request="cpu=0.3,memory=1024Mi" \
+  --chief_resource_limit="cpu=1,memory=2048Mi" \
+  --evaluator_resource_request="cpu=0.3,memory=1024Mi" \
+  --evaluator_resource_limit="cpu=1,memory=2048Mi" \
+  --job_name=test-deepfm-estimator \
+  --distribution_strategy=ParameterServerStrategy \
+  --need_tf_config=true \
+  --volume="host_path={criteo_data_path},mount_path=/data" \
+```
+
+`--image_name` is the image to launch the ElasticDL master which
+has nothing to do with the estimator model. The ElasticDL master is
+responsible for launching pod and assigning data shards to workers with
+elasticity and fault-tolerance.
+
+`{criteo_data_path}` is the absolute path of the `./data` with `criteo_sample.txt`.
+Here, the option `--volume="host_path={criteo_data_path},mount_path=/data"`
+bind mount it into the containers/pods.
+
+The option `--num_workers=1` tells the master to start a worker pod.
+The option `--num_ps=1` tells the master to start a ps pod.
+The option `--num_evaluator=1` tells the master to start an evaluator pod.
+
+And the master will start a chief worker for a TensorFlow estimator model by default.
+
+### Check Job Status
+
+After the job submission, we can run the command `kubectl get pods` to list
+related containers.
+
+```bash
+NAME                                     READY   STATUS    RESTARTS   AGE
+elasticdl-test-deepctr-estimator-master     1/1     Running   0          9s
+test-deepctr-estimator-edljob-chief-0       1/1     Running   0          6s
+test-deepctr-estimator-edljob-evaluator-0   0/1     Pending   0          6s
+test-deepctr-estimator-edljob-ps-0          1/1     Running   0          7s
+test-deepctr-estimator-edljob-worker-0      1/1     Running   0          6s
+```
+
+We can view the log of workers by `kubectl logs test-deepctr-estimator-edljob-chief-0`.
+
+```text
+INFO:tensorflow:global_step/sec: 4.84156
+INFO:tensorflow:global_step/sec: 4.84156
+INFO:tensorflow:Saving checkpoints for 203 into /data/ckpts/model.ckpt.
+INFO:tensorflow:Saving checkpoints for 203 into /data/ckpts/model.ckpt.
+INFO:tensorflow:global_step/sec: 7.05433
+INFO:tensorflow:global_step/sec: 7.05433
+```
diff --git a/docs/tutorials/elasticdl_deepctr.md b/docs/tutorials/elasticdl_deepctr_keras.md
similarity index 100%
rename from docs/tutorials/elasticdl_deepctr.md
rename to docs/tutorials/elasticdl_deepctr_keras.md
diff --git a/docs/tutorials/elasticdl_estimator.md b/docs/tutorials/elasticdl_estimator.md
new file mode 100644
index 000000000..bd875b076
--- /dev/null
+++ b/docs/tutorials/elasticdl_estimator.md
@@ -0,0 +1,223 @@
+# Train TensorFlow Estimator Models using ElasticDL on Personal Computer
+
+This document shows how to run an ElasticDL job to train a tf.estimator
+model using iris dataset on Minikube.
+
+## Prerequisites
+
+1. Install Minikube, preferably >= v1.11.0, following the installation
+   [guide](https://kubernetes.io/docs/tasks/tools/install-minikube).  Minikube
+   runs a single-node Kubernetes cluster in a virtual machine on your personal
+   computer.
+
+1. Install Docker CE, preferably >= 18.x, following the
+   [guide](https://docs.docker.com/docker-for-mac/install/) for building Docker
+   images containing user-defined models and the ElasticDL framework.
+
+1. Install Python, preferably >= 3.6, because the ElasticDL command-line tool is
+   in Python.
+
+## Models
+
+Among all machine learning toolkits that ElasticDL can work with, TensorFlow is
+the most tested and used.  In this tutorial, we use a model from the [model
+zoo](https://github.com/sql-machine-learning/elasticdl/tree/develop/model_zoo)
+directory.  This model is defined using TensorFlow estimator API.
+
+## Datasets
+
+We use the [iris](https://archive.ics.uci.edu/ml/machine-learning-databases/iris/iris.data)
+dataset in this tutorial.
+
+```bash
+mkdir ./data
+wget https://archive.ics.uci.edu/ml/machine-learning-databases/iris/iris.data -O ./data/iris.data
+```
+
+## The Kubernetes Cluster
+
+The following command starts a Kubernetes cluster locally using Minikube.  It
+uses [VirtualBox](https://www.virtualbox.org/), a hypervisor coming with
+macOS, to create the virtual machine cluster.
+
+```bash
+minikube start --vm-driver=virtualbox \
+  --cpus 2 --memory 6144 --disk-size=50gb 
+eval $(minikube docker-env)
+```
+
+The command `minikube docker-env` returns a set of Bash environment variable
+to configure your local environment to re-use the Docker daemon inside
+the Minikube instance.
+
+The following command is necessary to enable
+[RBAC](https://kubernetes.io/docs/reference/access-authn-authz/rbac/) of
+Kubernetes.
+
+```bash
+kubectl apply -f \
+  https://raw.githubusercontent.com/sql-machine-learning/elasticdl/develop/elasticdl/manifests/elasticdl-rbac.yaml
+```
+
+If you happen to live in a region where `raw.githubusercontent.com` is banned,
+you might want to Git clone the above repository to get the YAML file.
+
+## Install ElasticDL Client Tool
+
+The following command installs the command line tool `elasticdl`, which talks to
+the Kubernetes cluster and operates ElasticDL jobs.
+
+```bash
+pip install elasticdl_client
+```
+
+## Build the Docker Image with Model Definition
+
+Kubernetes runs Docker containers, so we need to put user-defined models,
+the ElasticDL api package and all dependencies into a Docker image.
+
+In this tutorial, we use a predefined model in the ElasticDL repository.  To
+retrieve the source code, please run the following command.
+
+```bash
+git clone https://github.com/sql-machine-learning/elasticdl
+```
+
+The estimator model definition is in directory [elasticdl/model_zoo/iris](https://github.com/sql-machine-learning/elasticdl/tree/develop/model_zoo/iris).
+
+We build the image based on tensorflow:1.13.2 and the dockerfile is
+
+```text
+FROM tensorflow/tensorflow:1.13.2-py3 as base
+
+RUN pip install elasticdl_api
+
+COPY ./model_zoo model_zoo
+```
+
+Then, we use docker to build the image
+
+```bash
+docker build -t elasticdl:iris_estimator -f ${iris_dockerfile} .
+```
+
+## Submit the Training Job
+
+The following command submits a training job:
+
+```bash
+elasticdl train \
+  --image_name=elasticdl:1.0.0 \
+  --worker_image=elasticdl:iris_estimator \
+  --ps_image=elasticdl:iris_estimator \
+  --job_command="python -m model_zoo.iris.dnn_estimator" \
+  --master_resource_request="cpu=0.2,memory=1024Mi" \
+  --master_resource_limit="cpu=1,memory=2048Mi" \
+  --num_ps=1 \
+  --ps_resource_request="cpu=0.2,memory=1024Mi" \
+  --ps_resource_limit="cpu=1,memory=2048Mi" \
+  --num_workers=1 \
+  --worker_resource_request="cpu=0.3,memory=1024Mi" \
+  --worker_resource_limit="cpu=1,memory=2048Mi" \
+  --chief_resource_request="cpu=0.3,memory=1024Mi" \
+  --chief_resource_limit="cpu=1,memory=2048Mi" \
+  --num_evaluator=1 \
+  --evaluator_resource_request="cpu=0.3,memory=1024Mi" \
+  --evaluator_resource_limit="cpu=1,memory=2048Mi" \
+  --job_name=test-iris-estimator \
+  --image_pull_policy=Never \
+  --distribution_strategy=ParameterServerStrategy \
+  --need_tf_config=true \
+  --volume="host_path={iris_data_dir},mount_path=/data" \
+```
+
+`--image_name` is the image to launch the ElasticDL master which
+has nothing to do with the estimator model. The ElasticDL master is
+responsible for launching pod and assigning data shards to workers with
+elasticity and fault-tolerance.
+
+`{iris_data_dir}` is the absolute path of the `./data` with `iris.data`.
+Here, the option `--volume="host_path={iris_data_dir},mount_path=/data"`
+bind mount it into the containers/pods.
+
+The option `--num_workers=1` tells the master to start a worker pod.
+The option `--num_ps=1` tells the master to start a ps pod.
+The option `--num_evaluator` tells the master to start an evaluator pod.
+
+And the master will start a chief worker for a TensorFlow estiamtor model by default.
+
+### Check Job Status
+
+After the job submission, we can run the command `kubectl get pods` to list
+related containers.
+
+```bash
+NAME                                     READY   STATUS    RESTARTS   AGE
+elasticdl-test-iris-estimator-master     1/1     Running   0          9s
+test-iris-estimator-edljob-chief-0       1/1     Running   0          6s
+test-iris-estimator-edljob-evaluator-0   0/1     Pending   0          6s
+test-iris-estimator-edljob-ps-0          1/1     Running   0          7s
+test-iris-estimator-edljob-worker-0      1/1     Running   0          6s
+```
+
+## Train an Estimator Model Using ElasticDL with Your Dataset
+
+You only need to modify your `input_fn` with ElasticDL DataShardService.
+The DataShardService will split the sample indices into ranges and assign
+those ranges to workers. The worker only need to read samples by indices
+in those ranges.
+
+1. Create a DataShardService.
+
+```python
+from elasticai_api.common.data_shard_service import build_data_shard_service
+
+training_data_shard_svc = build_data_shard_service(
+        batch_size=batch_size,
+        num_epochs=100,
+        dataset_size=len(rows),
+        num_minibatches_per_shard=1,
+        dataset_name="iris_training_data",
+    )
+```
+
+- batch_size: Batch size of each step.
+- num_epochs: The number of epochs.
+- dataset_size: The total number of samples in the dataset.
+- num_minibatches_per_shard: The number of batches in each shard.
+  The number of samples in each shard is
+  `batch_size * num_minibatches_per_shard`
+- dataset_name: The name of dataset.
+
+2. Create a generator by reading samples with shards.
+
+The `shard.start` and `shard.end` is the start index
+and end index of samples in those shard. You can read
+samples by the two indices like:
+
+```python
+def train_generator(shard_service):
+    while True:
+        # Read samples by the range of the shard from
+        # the data shard serice.
+        shard = shard_service.fetch_shard()
+        if not shard:
+            break
+        for i in range(shard.start, shard.end):
+            label = CATEGORY_CODE[rows[i][-1]]
+            yield rows[i][0:-1], [label]
+```
+
+3. Create a session hook to report shard
+
+```python
+from elasticai_api.tensorflow.hooks import ElasticDataShardReportHook
+
+hooks = [
+    ElasticDataShardReportHook(training_data_shard_svc),
+]
+train_spec = tf.estimator.TrainSpec(input_fn=train_input_fn, hooks=hooks)
+```
+
+After 3 steps, you can train your estimator models using ElasticDL
+in data parallel mode.
diff --git a/docs/tutorials/elasticdl_torch.md b/docs/tutorials/elasticdl_torch.md
new file mode 100644
index 000000000..5d6d1a068
--- /dev/null
+++ b/docs/tutorials/elasticdl_torch.md
@@ -0,0 +1,255 @@
+# Train PyTorch Models using ElasticDL on Personal Computer
+
+This document shows how to run an ElasticDL AllReduce job to train a PyTorch
+model using MNIST dataset on Minikube.
+
+## Prerequisites
+
+1. Install Minikube, preferably >= v1.11.0, following the installation
+   [guide](https://kubernetes.io/docs/tasks/tools/install-minikube).  Minikube
+   runs a single-node Kubernetes cluster in a virtual machine on your personal
+   computer.
+
+1. Install Docker CE, preferably >= 18.x, following the
+   [guide](https://docs.docker.com/docker-for-mac/install/) for building Docker
+   images containing user-defined models and the ElasticDL framework.
+
+1. Install Python, preferably >= 3.6, because the ElasticDL command-line tool is
+   in Python.
+
+## Models
+
+In this tutorial, we use the model defined in the [model
+zoo](https://github.com/sql-machine-learning/elasticdl/tree/develop/model_zoo/mnist/mnist_pytorch.py)
+directory.  This model is defined using PyTorch API.
+
+## Datasets
+
+We use the [MINST](https://www.kaggle.com/jidhumohan/mnist-png/download)
+dataset in this tutorial. After downloading the dataset, we should
+unzip it into a local directory.
+
+## The Kubernetes Cluster
+
+The following command starts a Kubernetes cluster locally using Minikube.  It
+uses [VirtualBox](https://www.virtualbox.org/), a hypervisor coming with
+macOS, to create the virtual machine cluster.
+
+```bash
+minikube start --vm-driver=virtualbox \
+  --cpus 2 --memory 6144 --disk-size=50gb 
+eval $(minikube docker-env)
+```
+
+The command `minikube docker-env` returns a set of Bash environment variable
+to configure your local environment to re-use the Docker daemon inside
+the Minikube instance.
+
+The following command is necessary to enable
+[RBAC](https://kubernetes.io/docs/reference/access-authn-authz/rbac/) of
+Kubernetes.
+
+```bash
+kubectl apply -f \
+  https://raw.githubusercontent.com/sql-machine-learning/elasticdl/develop/elasticdl/manifests/elasticdl-rbac.yaml
+```
+
+If you happen to live in a region where `raw.githubusercontent.com` is banned,
+you might want to Git clone the above repository to get the YAML file.
+
+## Install ElasticDL Client Tool
+
+The following command installs the command line tool `elasticdl`, which talks to
+the Kubernetes cluster and operates ElasticDL jobs.
+
+```bash
+pip install elasticdl_client
+```
+
+## Build the Docker Image with Model Definition
+
+Kubernetes runs Docker containers, so we need to put user-defined models,
+the ElasticDL api package and all dependencies into a Docker image.
+
+In this tutorial, we use a predefined model in the ElasticDL repository.  To
+retrieve the source code, please run the following command.
+
+```bash
+git clone https://github.com/sql-machine-learning/elasticdl
+```
+
+The model definition is in directory [elasticdl/model_zoo/mnist/mnist_pytorch.py](https://github.com/sql-machine-learning/elasticdl/blob/develop/model_zoo/mnist/mnist_pytorch.py).
+
+We build the image based on `horovod/horovod:0.20.0-tf2.3.0-torch1.6.0-mxnet1.5.0-py3.7-cpu`
+and the dockerfile is
+
+```text
+FROM horovod/horovod:0.20.0-tf2.3.0-torch1.6.0-mxnet1.5.0-py3.7-cpu as base
+
+RUN pip install opencv-python
+RUN apt update
+RUN apt install -y libgl1-mesa-glx libglib2.0-dev
+
+RUN HOROVOD_WITHOUT_MPI=1 \
+    HOROVOD_WITH_GLOO=1 \
+    HOROVOD_WITHOUT_MXNET=1 \
+    HOROVOD_WITH_TENSORFLOW=1 \
+    HOROVOD_WITH_PYTORCH=1  \
+    pip install horovod==0.21.0
+
+RUN pip install elasticdl_api
+
+COPY ./model_zoo model_zoo
+ENV PYTHONUNBUFFERED 0
+```
+
+Then, we use docker to build the image
+
+```bash
+docker build -t elasticdl:mnist_pytorch -f ${mnist_dockerfile} .
+```
+
+## Submit the Training Job
+
+The following command submits a training job:
+
+```bash
+elasticdl train \
+  --image_name=elasticdl/elasticdl:v1.0.0 \
+  --worker_image=elasticdl:mnist_pytorch \
+  --job_command="python -m model_zoo.mnist.mnist_pytorch --batch_size 64 --num_epochs 1 --training_data=/data/mnist_png/training --validation_data=/data/mnist_png/testing" \
+  --num_minibatches_per_task=2 \
+  --num_workers=2 \
+  --master_resource_request="cpu=0.2,memory=1024Mi" \
+  --master_resource_limit="cpu=1,memory=2048Mi" \
+  --worker_resource_request="cpu=0.3,memory=1024Mi" \
+  --worker_resource_limit="cpu=1,memory=2048Mi" \
+  --envs="USE_TORCH=true,HOROVOD_GLOO_TIMEOUT_SECONDS=60,PYTHONUNBUFFERED=true" \
+  --job_name=test-mnist-allreduce \
+  --image_pull_policy=Never \
+  --volume="host_path={mnist_data_dir},mount_path=/data" \
+  --distribution_strategy=AllreduceStrategy 
+```
+
+`--image_name` is the image to launch the ElasticDL master which
+has nothing to do with the estimator model. The ElasticDL master is
+responsible for launching pod and assigning data shards to workers with
+elasticity and fault-tolerance.
+
+`{mnist_data_dir}` is the absolute path of the `./data` with the directory of
+`mnist_png`. Here, the option `--volume="host_path={mnist_data_dir},mount_path=/data"`
+bind mount it into the containers/pods.
+
+The option `--num_workers=2` tells the master to start 2 worker pods.
+
+### Check Job Status
+
+After the job submission, we can run the command `kubectl get pods` to list
+related containers.
+
+```bash
+NAME                                    READY   STATUS    RESTARTS   AGE
+elasticdl-test-mnist-allreduce-master   1/1     Running   0          7s
+test-mnist-allreduce-edljob-worker-0    1/1     Running   0          5s
+test-mnist-allreduce-edljob-worker-1    1/1     Running   0          5s
+```
+
+## Train an PyTorch Model Using ElasticDL with Your Dataset
+
+In order to support fault-tolerance and elasticity with ElasticDL, you only
+need to create a custom dataset and wrap the function to perform forward and
+backward computation using ElasticDL APIs.
+
+### Create a Dataset With the RecordIndexService of ElasticDL
+
+ElasticDL can split the total dataset into shards and assign those
+shards to workers. If some workers fail, ElasticDL can re-assign
+shards of failed workers to other running workers. We can get sample
+indices in those shards by `RecordIndexService`. We can create a
+dataset to read images by indices from the `RecordIndexService`.
+
+```python
+class ElasticDataset(Dataset):
+    def __init__(self, images, data_shard_service=None):
+        """The dataset supports elastic training.
+
+        Args:
+            images: A list with tuples like (image_path, label_index).
+            For example, we can use `torchvision.datasets.ImageFolder`
+            to get the list.
+            data_shard_service: If we want to use elastic training, we
+            need to use the `data_shard_service` of the elastic controller
+            in elasticai_api.
+        """
+        self.data_shard_service = data_shard_service
+        self._images = images
+
+    def __len__(self):
+        if self.data_shard_service:
+            # Set the maxsize because the size of dataset is not fixed
+            # when using dynamic sharding
+            return sys.maxsize
+        else:
+            return len(self._images)
+
+    def __getitem__(self, index):
+        if self.data_shard_service:
+            index = self.data_shard_service.fetch_record_index()
+            return self.read_image(index)
+        else:
+            return self.read_image(index)
+
+    def read_image(self, index):
+        image_path, label = self._images[index]
+        image = cv2.imread(image_path)
+        image = np.array(image / 255.0, np.float32)
+        image = image.reshape(3, 28, 28)
+        return image, label
+
+
+if __name__ == "__main__":
+    ...
+    data_shard_service = RecordIndexService(
+        batch_size=args.batch_size,
+        dataset_size=len(train_data.imgs),
+        num_epochs=args.num_epochs,
+        shuffle=True,
+        dataset_name="mnist_training_data",
+    )
+    train_dataset = ElasticDataset(train_data.imgs, data_shard_service)
+    ...
+```
+
+### Create an ElasticDL Controller to Wrap the Forward and Backward Computation
+
+In ElasticDL AllReduce, we need to create a `PyTorchAllReduceController`
+of ElasticDL. At the begining, The controller can broadcast the model and
+optimizer. If some workers fail, the controller will re-initialize Horovod
+using a new world. After creating the controller, we should wrap the function
+to perform the forward and backward computation by `elastic_run`.
+
+```python
+def train_one_batch(model, optimizer, data, target):
+    optimizer.zero_grad()
+    output = model(data)
+    loss = F.nll_loss(output, target)
+    loss.backward()
+    optimizer.step()
+    return loss
+
+
+if __name__ == "__main__":
+    ...
+    model = ...
+    optimizer = ...
+    allreduce_controller = PyTorchAllReduceController(data_shard_service)
+    allreduce_controller.set_broadcast_model(model)
+    allreduce_controller.set_broadcast_optimizer(optimizer)
+    # Use the elastic function to wrap the training function with a batch.
+    elastic_train_one_batch = allreduce_controller.elastic_run(train_one_batch)
+
+    with allreduce_controller.scope():
+        for batch_idx, (data, target) in enumerate(train_loader):
+            loss = elastic_train_one_batch(model, optimizer, data, target)
+    ...
+```
\ No newline at end of file
diff --git a/elasticai_api/tensorflow/hooks.py b/elasticai_api/tensorflow/hooks.py
new file mode 100644
index 000000000..0f933ed02
--- /dev/null
+++ b/elasticai_api/tensorflow/hooks.py
@@ -0,0 +1,27 @@
+# Copyright 2021 The ElasticDL Authors. All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import tensorflow as tf
+
+from elasticai_api.util.log_utils import default_logger as logger
+
+
+class ElasticDataShardReportHook(tf.train.SessionRunHook):
+    def __init__(self, data_shard_service) -> None:
+        self._data_shard_service = data_shard_service
+
+    def after_run(self, run_context, run_values):
+        try:
+            self._data_shard_service.report_batch_done()
+        except Exception as ex:
+            logger.error("elastic_ai: report batch done failed: %s", ex)
diff --git a/elasticdl/python/tests/data_reader_test.py b/elasticdl/python/tests/data_reader_test.py
index b9631cde2..9869716b2 100644
--- a/elasticdl/python/tests/data_reader_test.py
+++ b/elasticdl/python/tests/data_reader_test.py
@@ -63,12 +63,13 @@ def test_recordio_data_reader(self):
                 )
             )
             self.assertEqual(len(records), num_records)
+            fixed_len_feature = tf.io.FixedLenFeature([1], tf.float32)
             for record in records:
                 parsed_record = tf.io.parse_single_example(
                     record,
                     {
-                        "x": tf.io.FixedLenFeature([1], tf.float32),
-                        "y": tf.io.FixedLenFeature([1], tf.float32),
+                        "x": fixed_len_feature,
+                        "y": fixed_len_feature,
                     },
                 )
                 for k, v in parsed_record.items():
diff --git a/elasticdl/requirements.txt b/elasticdl/requirements.txt
index 25cce98d7..401388ae3 100644
--- a/elasticdl/requirements.txt
+++ b/elasticdl/requirements.txt
@@ -3,6 +3,6 @@ docker==4.2.1
 pyrecordio>=0.0.6
 Cython
 odps
-tensorflow==2.1.2
+tensorflow==2.5.2
 psutil
 cloudpickle
diff --git a/elasticdl_client/requirements.txt b/elasticdl_client/requirements.txt
index c939e8f3b..339126a24 100644
--- a/elasticdl_client/requirements.txt
+++ b/elasticdl_client/requirements.txt
@@ -1,3 +1,3 @@
 kubernetes==12.0.0
 docker==4.2.1
-Jinja2==2.11.2
+Jinja2==2.11.3
diff --git a/elasticdl_preprocessing/requirements.txt b/elasticdl_preprocessing/requirements.txt
index 0169fb67c..d918bd039 100644
--- a/elasticdl_preprocessing/requirements.txt
+++ b/elasticdl_preprocessing/requirements.txt
@@ -1 +1 @@
-tensorflow==2.1.2
\ No newline at end of file
+tensorflow==2.9.3
\ No newline at end of file
diff --git a/model_zoo/deepctr/deepfm_estimator.py b/model_zoo/deepctr/deepfm_estimator.py
new file mode 100644
index 000000000..8d668d215
--- /dev/null
+++ b/model_zoo/deepctr/deepfm_estimator.py
@@ -0,0 +1,172 @@
+# Copyright 2021 The ElasticDL Authors. All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import csv
+import os
+
+import tensorflow as tf
+from deepctr.estimator.models import DeepFMEstimator
+
+from elasticai_api.common.data_shard_service import DataShardService
+from elasticai_api.tensorflow.hooks import ElasticDataShardReportHook
+
+tf.logging.set_verbosity(tf.logging.INFO)
+
+
+def read_csv(file_path):
+    rows = []
+    with open(file_path) as csvfile:
+        spamreader = csv.reader(csvfile)
+        for i, row in enumerate(spamreader):
+            if i > 0:
+                row_values = []
+                row_values.append(int(row[0]))
+                for i in range(1, 14):
+                    value = row[i] if row[i] else 0
+                    row_values.append(float(value))
+                row_values.extend(row[14:])
+                rows.append(row_values)
+    return rows
+
+
+def train_generator(data_path, shard_service):
+    rows = read_csv(data_path)
+    while True:
+        # Read samples by the range of the shard from
+        # the data shard serice.
+        shard = shard_service.fetch_shard()
+        if not shard:
+            break
+        for i in range(shard.start, shard.end):
+            yield tuple(rows[i])
+
+
+def eval_generator(data_path):
+    rows = read_csv(data_path)
+    for row in rows:
+        yield tuple(row)
+
+
+def input_fn(sample_generator, batch_size, dense_features, sparse_features):
+    output_types = tuple(
+        [tf.int32]
+        + [tf.float32 for i in dense_features]
+        + [tf.string for i in sparse_features]
+    )
+    dataset = tf.data.Dataset.from_generator(
+        sample_generator, output_types=output_types,
+    )
+    dataset = dataset.shuffle(100).batch(batch_size)
+    values = dataset.make_one_shot_iterator().get_next()
+
+    label_value = values[0]
+    feature_values = {}
+    feature_index = 1
+    for feature in dense_features:
+        feature_values[feature] = values[feature_index]
+        feature_index += 1
+
+    for feature in sparse_features:
+        feature_values[feature] = values[feature_index]
+        feature_index += 1
+    return feature_values, label_value
+
+
+def arg_parser():
+    parser = argparse.ArgumentParser(description="Process training parameters")
+    parser.add_argument("--training_data", type=str, required=True)
+    parser.add_argument(
+        "--validation_data", type=str, default="", required=False
+    )
+    return parser
+
+
+if __name__ == "__main__":
+    parser = arg_parser()
+    args = parser.parse_args()
+
+    training_data = args.training_data
+    validation_data = args.validation_data
+
+    model_dir = "/data/ckpts/"
+    os.makedirs(model_dir, exist_ok=True)
+
+    sparse_features = ["C" + str(i) for i in range(1, 27)]
+    dense_features = ["I" + str(i) for i in range(1, 14)]
+
+    dnn_feature_columns = []
+    linear_feature_columns = []
+
+    for i, feat in enumerate(sparse_features):
+        dnn_feature_columns.append(
+            tf.feature_column.embedding_column(
+                tf.feature_column.categorical_column_with_hash_bucket(
+                    feat, 1000
+                ),
+                4,
+            )
+        )
+        linear_feature_columns.append(
+            tf.feature_column.categorical_column_with_hash_bucket(feat, 1000)
+        )
+    for feat in dense_features:
+        dnn_feature_columns.append(tf.feature_column.numeric_column(feat))
+        linear_feature_columns.append(tf.feature_column.numeric_column(feat))
+
+    batch_size = 64
+
+    config = tf.estimator.RunConfig(
+        model_dir=model_dir, save_checkpoints_steps=100, keep_checkpoint_max=3
+    )
+    model = DeepFMEstimator(
+        linear_feature_columns,
+        dnn_feature_columns,
+        task="binary",
+        config=config,
+    )
+
+    # Create a data shard service which can split the dataset
+    # into shards.
+    rows = read_csv(training_data)
+    training_data_shard_svc = DataShardService(
+        batch_size=batch_size,
+        num_epochs=100,
+        dataset_size=len(rows),
+        num_minibatches_per_shard=1,
+        dataset_name="iris_training_data",
+    )
+
+    def train_input_fn():
+        return input_fn(
+            lambda: train_generator(training_data, training_data_shard_svc),
+            batch_size,
+            dense_features,
+            sparse_features,
+        )
+
+    def eval_input_fn():
+        return input_fn(
+            lambda: eval_generator(validation_data),
+            batch_size,
+            dense_features,
+            sparse_features,
+        )
+
+    hooks = [
+        ElasticDataShardReportHook(training_data_shard_svc),
+    ]
+    train_spec = tf.estimator.TrainSpec(input_fn=train_input_fn, hooks=hooks)
+    eval_spec = tf.estimator.EvalSpec(input_fn=eval_input_fn)
+
+    tf.estimator.train_and_evaluate(model, train_spec, eval_spec)
diff --git a/model_zoo/iris/__init__.py b/model_zoo/iris/__init__.py
new file mode 100644
index 000000000..87de1e085
--- /dev/null
+++ b/model_zoo/iris/__init__.py
@@ -0,0 +1,12 @@
+# Copyright 2021 The ElasticDL Authors. All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
diff --git a/model_zoo/iris/dnn_estimator.py b/model_zoo/iris/dnn_estimator.py
new file mode 100644
index 000000000..614e9b077
--- /dev/null
+++ b/model_zoo/iris/dnn_estimator.py
@@ -0,0 +1,158 @@
+# Copyright 2021 The ElasticDL Authors. All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import csv
+import os
+
+import tensorflow as tf
+
+from elasticai_api.common.data_shard_service import build_data_shard_service
+from elasticai_api.tensorflow.hooks import ElasticDataShardReportHook
+
+tf.logging.set_verbosity(tf.logging.INFO)
+
+CATEGORY_CODE = {"Iris-setosa": 0, "Iris-versicolor": 1, "Iris-virginica": 2}
+DATASET_DIR = "/data/iris.data"
+
+
+def read_csv(file_path):
+    rows = []
+    with open(file_path) as csvfile:
+        csv_reader = csv.reader(csvfile)
+        for row in csv_reader:
+            rows.append(row)
+    return rows
+
+
+def model_fn(features, labels, mode, params):
+    net = tf.feature_column.input_layer(features, params["feature_columns"])
+
+    for units in params["hidden_units"]:
+        net = tf.layers.dense(net, units=units, activation=tf.nn.relu)
+    logits = tf.layers.dense(net, params["n_classes"], activation=None)
+
+    predicted_classes = tf.argmax(logits, 1)
+    if mode == tf.estimator.ModeKeys.PREDICT:
+        predictions = {
+            "classes": predicted_classes[:, tf.newaxis],
+            "probs": tf.nn.softmax(logits),
+            "logits": logits,
+        }
+        export_outputs = {
+            "prediction": tf.estimator.export.PredictOutput(predictions)
+        }
+        return tf.estimator.EstimatorSpec(
+            mode, predictions=predictions, export_outputs=export_outputs
+        )
+
+    loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)
+    if mode == tf.estimator.ModeKeys.TRAIN:
+        optimizer = tf.train.AdagradOptimizer(learning_rate=0.1)
+        train_op = optimizer.minimize(
+            loss, global_step=tf.train.get_global_step()
+        )
+        logging_hook = tf.train.LoggingTensorHook(
+            {"loss": loss}, every_n_iter=10
+        )
+        return tf.estimator.EstimatorSpec(
+            mode, loss=loss, train_op=train_op, training_hooks=[logging_hook]
+        )
+
+    accuracy = tf.metrics.accuracy(
+        labels=labels, predictions=predicted_classes, name="acc"
+    )
+    metrics = {"accuracy": accuracy}
+    if mode == tf.estimator.ModeKeys.EVAL:
+        return tf.estimator.EstimatorSpec(
+            mode, loss=loss, eval_metric_ops=metrics
+        )
+
+
+def train_generator(shard_service):
+    rows = read_csv(DATASET_DIR)
+    while True:
+        # Read samples by the range of the shard from
+        # the data shard serice.
+        shard = shard_service.fetch_shard()
+        if not shard:
+            break
+        for i in range(shard.start, shard.end):
+            label = CATEGORY_CODE[rows[i][-1]]
+            yield rows[i][0:-1], [label]
+
+
+def eval_generator():
+    rows = read_csv(DATASET_DIR)
+    for row in rows:
+        label = CATEGORY_CODE[row[-1]]
+        yield row[0:-1], [label]
+
+
+def input_fn(sample_generator, batch_size):
+    dataset = tf.data.Dataset.from_generator(
+        sample_generator,
+        output_types=(tf.float32, tf.int32),
+        output_shapes=(4, 1),
+    )
+    dataset = dataset.shuffle(100).batch(batch_size)
+    feature_values, label_values = dataset.make_one_shot_iterator().get_next()
+    features = {"x": feature_values}
+    return features, label_values
+
+
+if __name__ == "__main__":
+    model_dir = "/data/ckpts/"
+    batch_size = 64
+    feature_columns = [
+        tf.feature_column.numeric_column(key="x", shape=(4,), dtype=tf.float32)
+    ]
+    os.makedirs(model_dir, exist_ok=True)
+
+    config = tf.estimator.RunConfig(
+        model_dir=model_dir, save_checkpoints_steps=300, keep_checkpoint_max=5
+    )
+    classifier = tf.estimator.Estimator(
+        model_fn=model_fn,
+        config=config,
+        params={
+            "hidden_units": [8, 4],
+            "n_classes": 3,
+            "feature_columns": feature_columns,
+        },
+    )
+
+    # Create a data shard service which can split the dataset
+    # into shards.
+    rows = read_csv(DATASET_DIR)
+    training_data_shard_svc = build_data_shard_service(
+        batch_size=batch_size,
+        num_epochs=100,
+        dataset_size=len(rows),
+        num_minibatches_per_shard=1,
+        dataset_name="iris_training_data",
+    )
+
+    # Add a hook to report the shard done so that the data
+    # shard service will not reassign the shard to other workers.
+    hooks = [ElasticDataShardReportHook(training_data_shard_svc)]
+
+    def train_input_fn():
+        return input_fn(
+            lambda: train_generator(training_data_shard_svc), batch_size
+        )
+
+    train_spec = tf.estimator.TrainSpec(input_fn=train_input_fn, hooks=hooks,)
+    eval_spec = tf.estimator.EvalSpec(
+        input_fn=lambda: input_fn(eval_generator, batch_size)
+    )
+    tf.estimator.train_and_evaluate(classifier, train_spec, eval_spec)
diff --git a/model_zoo/mnist/mnist_pytorch.py b/model_zoo/mnist/mnist_pytorch.py
index 7dc3ccaad..f2f470141 100644
--- a/model_zoo/mnist/mnist_pytorch.py
+++ b/model_zoo/mnist/mnist_pytorch.py
@@ -11,31 +11,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""
-Download the mnist dataset from
-https://s3.amazonaws.com/fast-ai-imageclas/mnist_png.tgz
-and then untar it into ${data_store_dir}. Using minikube, we can use the
-following command to submit a training job with these codes.
-
-elasticdl train \
-  --image_name=elasticdl:pt_mnist_allreduce  \
-  --job_command="python -m model_zoo.mnist.mnist_pytorch_custom_dataset \
-      --training_data=/local_data/mnist_png/training \
-      --validation_data=/local_data/mnist_png/testing" \
-  --num_minibatches_per_task=2 \
-  --num_workers=2 \
-  --worker_pod_priority=0.5 \
-  --master_resource_request="cpu=0.2,memory=1024Mi" \
-  --master_resource_limit="cpu=1,memory=2048Mi" \
-  --worker_resource_request="cpu=0.3,memory=1024Mi" \
-  --worker_resource_limit="cpu=1,memory=2048Mi" \
-  --envs="USE_TORCH=true,HOROVOD_GLOO_TIMEOUT_SECONDS=60,PYTHONUNBUFFERED=0" \
-  --job_name=test-mnist-allreduce \
-  --image_pull_policy=Never \
-  --volume="host_path=${data_store_dir},mount_path=/local_data" \
-  --distribution_strategy=AllreduceStrategy \
-"""
-
 import argparse
 import sys
 
@@ -49,7 +24,8 @@
 from torch.optim.lr_scheduler import StepLR
 from torch.utils.data import DataLoader, Dataset
 
-from elasticai_api.pytorch.controller import create_elastic_controller
+from elasticai_api.common.data_shard_service import RecordIndexService
+from elasticai_api.pytorch.controller import PyTorchAllReduceController
 from elasticai_api.pytorch.optimizer import DistributedOptimizer
 
 
@@ -131,15 +107,14 @@ def train(args):
     train_data = torchvision.datasets.ImageFolder(args.training_data)
     test_data = torchvision.datasets.ImageFolder(args.validation_data)
 
-    allreduce_controller = create_elastic_controller(
+    data_shard_service = RecordIndexService(
         batch_size=args.batch_size,
         dataset_size=len(train_data.imgs),
         num_epochs=args.num_epochs,
         shuffle=True,
+        dataset_name="mnist_training_data",
     )
-    train_dataset = ElasticDataset(
-        train_data.imgs, allreduce_controller.data_shard_service
-    )
+    train_dataset = ElasticDataset(train_data.imgs, data_shard_service)
     train_loader = DataLoader(
         dataset=train_dataset, batch_size=args.batch_size, num_workers=2
     )
@@ -155,6 +130,7 @@ def train(args):
     scheduler = StepLR(optimizer, step_size=1, gamma=0.5)
 
     # Set the model and optimizer to broadcast.
+    allreduce_controller = PyTorchAllReduceController(data_shard_service)
     allreduce_controller.set_broadcast_model(model)
     allreduce_controller.set_broadcast_optimizer(optimizer)
     epoch = 0
diff --git a/model_zoo/mnist/mnist_tfv1_img.py b/model_zoo/mnist/mnist_tfv1_img.py
new file mode 100644
index 000000000..2e77802ff
--- /dev/null
+++ b/model_zoo/mnist/mnist_tfv1_img.py
@@ -0,0 +1,208 @@
+# Copyright 2020 The ElasticDL Authors. All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Download the mnist dataset from
+https://s3.amazonaws.com/fast-ai-imageclas/mnist_png.tgz
+and then untar it into ${data_store_dir}. On minikube, we can use the
+following command to submit a training job with these codes.
+
+elasticdl train \
+  --image_name=elasticdl:pt_mnist_allreduce  \
+  --job_command="python -m model_zoo.mnist.mnist_tfv1_png \
+      --training_data=/local_data/mnist_png/training/"
+  --num_minibatches_per_task=2 \
+  --num_workers=2 \
+  --worker_pod_priority=0.5 \
+  --master_resource_request="cpu=0.2,memory=1024Mi" \
+  --master_resource_limit="cpu=1,memory=2048Mi" \
+  --worker_resource_request="cpu=0.3,memory=1024Mi" \
+  --worker_resource_limit="cpu=1,memory=2048Mi" \
+  --envs="PYTHONUNBUFFERED=0,HOROVOD_ELASTIC=1" \
+  --job_name=test-mnist-allreduce \
+  --image_pull_policy=Never \
+  --volume="host_path=${data_store_dir},mount_path=/local_data" \
+  --distribution_strategy=AllreduceStrategy \
+"""
+
+import argparse
+import os
+
+import cv2
+import numpy as np
+import tensorflow as tf
+
+from elasticai_api.tensorflow.controller import create_elastic_controller
+from elasticai_api.tensorflow.optimizer import (
+    AdjustBackwardPassesPerStepHook,
+    DistributedOptimizer,
+)
+from elasticdl.python.common.log_utils import default_logger as logger
+
+layers = tf.layers
+
+
+def get_samples_from_folder(folder_dir):
+    category_index = 0
+    samples = []
+    for category_name in sorted(os.listdir(folder_dir)):
+        category_dir = os.path.join(folder_dir, category_name)
+        if not os.path.isdir(category_dir):
+            continue
+        for img_file in os.listdir(category_dir):
+            img_dir = os.path.join(category_dir, img_file)
+            if os.path.isfile(img_dir) and img_dir.endswith("jpg"):
+                samples.append((img_dir, category_index))
+        category_index += 1
+    return samples
+
+
+def get_dataset_gen(data_shard_service, samples):
+    def gen():
+        while True:
+            # The data shard service will fetch the index of sample
+            # from the master of ElasticDL job.
+            index = data_shard_service.fetch_record_index()
+            if not index:
+                raise StopIteration("No data")
+            image_path, label = samples[index]
+            image = cv2.imread(image_path)
+            image = np.array(image / 255.0, np.float32)
+            yield image, np.array([label])
+
+    return gen
+
+
+def create_dataset(data_shard_service, samples):
+    gen = get_dataset_gen(data_shard_service, samples)
+    dataset = tf.data.Dataset.from_generator(gen, (tf.float32, tf.float32))
+    return dataset
+
+
+def conv_model(feature, target, mode):
+    """2-layer convolution model."""
+    # Convert the target to a one-hot tensor of shape (batch_size, 10) and
+    # with a on-value of 1 for each one-hot vector of length 10.
+    target = tf.one_hot(tf.cast(target, tf.int32), 10, 1, 0)
+
+    # Reshape feature to 4d tensor with 2nd and 3rd dimensions being
+    # image width and height final dimension being the number of color
+    # channels.
+    feature = tf.reshape(feature, [-1, 28, 28, 3])
+
+    # First conv layer will compute 32 features for each 5x5 patch
+    with tf.variable_scope("conv_layer1"):
+        h_conv1 = layers.conv2d(
+            feature,
+            32,
+            kernel_size=[5, 5],
+            activation=tf.nn.relu,
+            padding="SAME",
+        )
+        h_pool1 = tf.nn.max_pool(
+            h_conv1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding="SAME"
+        )
+
+    # Second conv layer will compute 64 features for each 5x5 patch.
+    with tf.variable_scope("conv_layer2"):
+        h_conv2 = layers.conv2d(
+            h_pool1,
+            64,
+            kernel_size=[5, 5],
+            activation=tf.nn.relu,
+            padding="SAME",
+        )
+        h_pool2 = tf.nn.max_pool(
+            h_conv2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding="SAME"
+        )
+        # reshape tensor into a batch of vectors
+        h_pool2_flat = tf.reshape(h_pool2, [-1, 7 * 7 * 64])
+
+    # Densely connected layer with 1024 neurons.
+    h_fc1 = layers.dropout(
+        layers.dense(h_pool2_flat, 1024, activation=tf.nn.relu),
+        rate=0.5,
+        training=mode == tf.estimator.ModeKeys.TRAIN,
+    )
+
+    # Compute logits (1 per class) and compute loss.
+    logits = layers.dense(h_fc1, 10, activation=None)
+    loss = tf.losses.softmax_cross_entropy(target, logits)
+
+    return tf.argmax(logits, 1), loss
+
+
+def train(args):
+    training_samples = get_samples_from_folder(args.training_data)
+    allreduce_controller = create_elastic_controller(
+        batch_size=args.batch_size,
+        num_epochs=args.num_epochs,
+        dataset_size=len(training_samples),
+        shuffle=True,
+    )
+
+    # Create the dataset with the elastic allreduce controller
+    # which will fetch data shards from the master of ElasticDL job.
+    dataset = create_dataset(
+        allreduce_controller.data_shard_service, training_samples
+    )
+    dataset = dataset.batch(args.batch_size).prefetch(1)
+    dataset_it = dataset.make_one_shot_iterator()
+    batch_x, batch_y = dataset_it.get_next()
+    batch_x = tf.cast(batch_x, tf.float32)
+
+    batch_y = tf.reshape(batch_y, (-1,))
+    predict, loss = conv_model(batch_x, batch_y, tf.estimator.ModeKeys.TRAIN)
+    optimizer = tf.train.GradientDescentOptimizer(0.1)
+
+    # Wrap the optimizer.
+    optimizer = DistributedOptimizer(optimizer, fixed_global_batch_size=True)
+    global_step = tf.train.get_or_create_global_step()
+    train_step = optimizer.minimize(loss, global_step=global_step)
+
+    # Use the elastic wrapper to wrap the function to train one batch
+    elastic_train_one_batch = allreduce_controller.elastic_run(train_one_batch)
+    hook = AdjustBackwardPassesPerStepHook(optimizer)
+    allreduce_controller.set_broadcast_variables(tf.global_variables())
+    with allreduce_controller.scope():
+        with tf.train.MonitoredTrainingSession(hooks=[hook]) as sess:
+            allreduce_controller.set_session(sess)
+            try:
+                while True:
+                    loss_value, step, _ = elastic_train_one_batch(
+                        sess, [loss, global_step, train_step]
+                    )
+                    logger.info(
+                        "global step = {}. loss: {}".format(step, loss_value)
+                    )
+            except tf.errors.OutOfRangeError:
+                print("end!")
+
+
+def train_one_batch(sess, run_tensors):
+    return sess.run(run_tensors)
+
+
+def arg_parser():
+    parser = argparse.ArgumentParser(description="Process training parameters")
+    parser.add_argument("--batch_size", type=int, default=64, required=False)
+    parser.add_argument("--num_epochs", type=int, default=1, required=False)
+    parser.add_argument("--training_data", type=str, required=True)
+    return parser
+
+
+if __name__ == "__main__":
+    parser = arg_parser()
+    args = parser.parse_args()
+    print(args)
+    train(args)
diff --git a/model_zoo/mnist/mnist_train_tfv1.py b/model_zoo/mnist/mnist_tfv1_recordio.py
similarity index 100%
rename from model_zoo/mnist/mnist_train_tfv1.py
rename to model_zoo/mnist/mnist_tfv1_recordio.py
