def device_policy(self, policy):
handle = self._handle
old = pywrap_tensorflow.TFE_ContextGetDevicePlacementPolicy(handle)
pywrap_tensorflow.TFE_ContextSetThreadLocalDevicePlacementPolicy(
handle, policy)
try:
yield
finally:
pywrap_tensorflow.TFE_ContextSetThreadLocalDevicePlacementPolicy(
handle, old)
def device_policy(self, policy):
if not self._context_handle:
self._initialize_handle_and_devices()
old = pywrap_tensorflow.TFE_ContextGetDevicePlacementPolicy(
self._context_handle)
pywrap_tensorflow.TFE_ContextSetThreadLocalDevicePlacementPolicy(
self._handle, policy)
try:
yield
finally:
pywrap_tensorflow.TFE_ContextSetThreadLocalDevicePlacementPolicy(
self._handle, old)
def device_policy(self, policy):
if policy is None:
policy = DEVICE_PLACEMENT_SILENT
if self._device_policy != policy:
self._device_policy = policy
# Only set the policy if the context has already been initialized
if self._context_handle is not None:
pywrap_tensorflow.TFE_ContextSetThreadLocalDevicePlacementPolicy(
self._handle, self._device_policy)