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)
Example #2
0
 def _convert(pb_str):
   m = device_attributes_pb2.DeviceAttributes()
   m.ParseFromString(pb_str)
   return m