Update to TF 2.0

Current code is not compatible with TF 2.0.