Add EPU device type. Allow passing string to initialize_system_for_tpu_embedding.
PiperOrigin-RevId: 334657063 Change-Id: I6c339b32fb118fd9d8ac93c55656fe6915484aae
This commit is contained in:
parent
9fe4c92fc7
commit
d3a378f966
@ -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"})
|
||||
|
||||
|
||||
# ==============================================================================
|
||||
|
Loading…
x
Reference in New Issue
Block a user