def _get_device_attrs(device_type, device_name=None, cc_major=None, cc_minor=None): if device_type == 'CPU': return device_attributes_pb2.DeviceAttributes(device_type='CPU') assert device_type == 'GPU', 'Invalid device type: %s' % (device_type,) if not device_name: return device_attributes_pb2.DeviceAttributes(device_type='GPU') physical_device_desc = ( 'device: 0, name: %s, pci bus id: 0:0:0.0' % device_name) if cc_major: physical_device_desc += ', compute capability: %d.%d' % (cc_major, cc_minor) return device_attributes_pb2.DeviceAttributes( device_type='GPU', physical_device_desc=physical_device_desc)
def _convert(pb_str): m = device_attributes_pb2.DeviceAttributes() m.ParseFromString(pb_str) return m