def testInvalidDeviceNumber(self): opts = tf_session.TF_NewSessionOptions() c_session = tf_session.TF_NewSession(ops.get_default_graph()._c_graph, opts) raw_device_list = tf_session.TF_SessionListDevices(c_session) size = tf_session.TF_DeviceListCount(raw_device_list) with self.assertRaises(errors.InvalidArgumentError): tf_session.TF_DeviceListMemoryBytes(raw_device_list, size) tf_session.TF_DeleteDeviceList(raw_device_list) tf_session.TF_CloseSession(c_session)
def testInvalidDeviceNumber(self): opts = tf_session.TF_NewSessionOptions() with errors.raise_exception_on_not_ok_status() as status: c_session = tf_session.TF_NewSession( ops.get_default_graph()._c_graph, opts, status) raw_device_list = tf_session.TF_SessionListDevices( c_session, status) size = tf_session.TF_DeviceListCount(raw_device_list) # Test that invalid device numbers return -1 rather than a Swig-wrapped # pointer. status_no_exception = c_api_util.ScopedTFStatus() memory = tf_session.TF_DeviceListMemoryBytes(raw_device_list, size, status_no_exception) self.assertEqual(memory, -1) tf_session.TF_DeleteDeviceList(raw_device_list) with errors.raise_exception_on_not_ok_status() as status: tf_session.TF_CloseSession(c_session, status)
def list_devices(session): with errors.raise_exception_on_not_ok_status() as status: if session._created_with_new_api: raw_device_list = tf_session.TF_SessionListDevices( session._session, status) else: raw_device_list = tf_session.TF_DeprecatedSessionListDevices( session._session, status) device_list = [] size = tf_session.TF_DeviceListCount(raw_device_list) for i in range(size): name = tf_session.TF_DeviceListName(raw_device_list, i, status) device_type = tf_session.TF_DeviceListType(raw_device_list, i, status) memory = 0 # tf_session.TF_DeviceListMemoryBytes(raw_device_list, i, status) device_list.append(_DeviceAttributes(name, device_type, memory)) tf_session.TF_DeleteDeviceList(raw_device_list) return device_list