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)