diff --git a/elasticdl/Makefile b/elasticdl/Makefile
index 981b1a0fc..446ad3de7 100644
--- a/elasticdl/Makefile
+++ b/elasticdl/Makefile
@@ -1,21 +1,21 @@
 PYTHON = python
-
-EDL_PROTO_FILE = elasticdl/proto/elasticdl.proto
-
 GO_FILE_DIR = elasticdl/pkg/proto
-GO_PB_FILE = $(GO_FILE_DIR)/elasticdl.pb.go
-
 CXX = g++
 CAPI_DIR = elasticdl/pkg/kernel/capi
 
-all: python_pb $(GO_PB_FILE) kernel
+all: python_pb go_pb kernel
 
-python_pb:
-	$(PYTHON) -m grpc_tools.protoc -I . $(EDL_PROTO_FILE) --python_out=. --grpc_python_out=.
+python_pb: elasticdl/proto/elasticdl.proto
+	$(PYTHON) -m grpc_tools.protoc -I${TF_PATH} -I. \
+	    elasticdl/proto/elasticdl.proto --python_out=. --grpc_python_out=.
 
-$(GO_PB_FILE): $(EDL_PROTO_FILE)
+go_pb: elasticdl/proto/elasticdl.proto
 	mkdir -p $(GO_FILE_DIR)
-	protoc -I . $< --go_out=plugins=grpc:.
+	protoc -I${TF_PATH} ${TF_PATH}/tensorflow/core/framework/types.proto --go_out=${GOPATH}/src
+	protoc -I${TF_PATH} ${TF_PATH}/tensorflow/core/framework/resource_handle.proto --go_out=${GOPATH}/src
+	protoc -I${TF_PATH} ${TF_PATH}/tensorflow/core/framework/tensor_shape.proto --go_out=${GOPATH}/src
+	protoc -I${TF_PATH} ${TF_PATH}/tensorflow/core/framework/tensor.proto --go_out=${GOPATH}/src
+	protoc -I${TF_PATH} -Ielasticdl/proto elasticdl/proto/elasticdl.proto --go_out=plugins=grpc:.
 
 kernel:
 	${CXX} -O3 -std=c++11 ${CAPI_DIR}/kernel_api.cc -c -o ${CAPI_DIR}/libkernel_api.o
diff --git a/elasticdl/docker/Dockerfile.dev b/elasticdl/docker/Dockerfile.dev
index 04a164392..384a91891 100644
--- a/elasticdl/docker/Dockerfile.dev
+++ b/elasticdl/docker/Dockerfile.dev
@@ -49,11 +49,11 @@ RUN /install-protobuf.bash && rm /install-protobuf.bash
 
 # Install elasticdl.org/elasticdl Go package
 ENV ELASTICDLPATH $GOPATH/src/elasticdl.org/elasticdl
+ENV TF_PATH /tmp/tensorflow
+RUN cd /tmp && git clone --depth=1 https://github.com/tensorflow/tensorflow.git
 COPY . /tmp/elasticdl
-RUN cd /tmp/elasticdl && make -f elasticdl/Makefile && \
-    mkdir -p $ELASTICDLPATH && \
+RUN cd /tmp/elasticdl && mkdir -p $ELASTICDLPATH && \
+    make -f elasticdl/Makefile && \
     cp -r /tmp/elasticdl/elasticdl/pkg $ELASTICDLPATH && \
     cd $GOPATH/src/elasticdl.org/elasticdl && \
-    export GO111MODULE=on && \
-    go mod init elasticdl.org/elasticdl && \
     go install ./... && cd / && rm -rf /tmp/elasticdl
diff --git a/elasticdl/docker/scripts/install-go.bash b/elasticdl/docker/scripts/install-go.bash
index 69d32c86c..9650bbc26 100755
--- a/elasticdl/docker/scripts/install-go.bash
+++ b/elasticdl/docker/scripts/install-go.bash
@@ -6,7 +6,6 @@ GO_MIRROR_URL=$1
 
 curl --silent ${GO_MIRROR_URL}/go1.13.4.linux-amd64.tar.gz | tar -C /usr/local -xzf -
 
-go env -w GO111MODULE=on
 go env -w GOPROXY=https://goproxy.io,direct
 
 go get github.com/golang/protobuf/protoc-gen-go
@@ -16,5 +15,6 @@ go get golang.org/x/tools/cmd/cover
 go get github.com/mattn/goveralls
 go get github.com/rakyll/gotest
 go get github.com/wangkuiyi/goyaccfmt
+go get github.com/stretchr/testify/assert
 
 cp $GOPATH/bin/* /usr/local/bin/
diff --git a/elasticdl/proto/elasticdl.proto b/elasticdl/proto/elasticdl.proto
index 33dd3e37e..316967bb6 100644
--- a/elasticdl/proto/elasticdl.proto
+++ b/elasticdl/proto/elasticdl.proto
@@ -1,6 +1,7 @@
 syntax = "proto3";
 
 import "google/protobuf/empty.proto";
+import "tensorflow/core/framework/tensor.proto";
 
 package proto;
 option go_package = "elasticdl/pkg/proto";