Add EPU device type. Allow passing string to initialize_system_for_tpu_embedding.

PiperOrigin-RevId: 334657063
Change-Id: I6c339b32fb118fd9d8ac93c55656fe6915484aae
This commit is contained in:
Bruce Fontaine 2020-09-30 12:40:59 -07:00 committed by TensorFlower Gardener
parent 9fe4c92fc7
commit d3a378f966

View File

@ -21,7 +21,8 @@ from __future__ import print_function
from tensorflow.python.util.tf_export import tf_export
_VALID_DEVICE_TYPES = frozenset({"CPU", "GPU", "TPU", "CUSTOM"})
# EPU represents for TPU embedding for now. Subject to change in future.
_VALID_DEVICE_TYPES = frozenset({"CPU", "GPU", "TPU", "CUSTOM", "EPU"})
# ==============================================================================