Update dtensor GCP example with latest code
Update dtensor GCP example with latest code

This is a mimic from my internal code in g3/experimental
tensorflow/dtensor-gcp-examples,tensorflow/dtensor-gcp-examples