def _split_and_init()

in graphlearn_torch/python/data/feature.py [0:0]

• 26 lines of code
• 7 McCabe index (conditional complexity)

  def _split_and_init(self):
    r""" Split cpu feature tensor and initialize GPU part and CPU part.
    """
    self._check_and_set_device()
    device_part, cpu_part = self._split(self.feature_tensor)

    if device_part.shape[0] > 0: # GPU part
      for group in self.device_group_list:
        block_size = device_part.shape[0] // group.size
        unified_tensor = UnifiedTensor(group.device_list[0], self.dtype)
        tensors, tensor_devices = [], []
        cur_pos = 0
        for idx, device in enumerate(group.device_list):
          if idx == group.size - 1:
            tensors.append(device_part[cur_pos:])
          else:
            tensors.append(device_part[cur_pos:cur_pos + block_size])
            cur_pos += block_size
          tensor_devices.append(device)
        unified_tensor.init_from(tensors, tensor_devices)
        self._unified_tensors[group.group_id] = unified_tensor

    if cpu_part.numel() > 0: # CPU part
      group_id = self._device2group[self.device]
      unified_tensor = self._unified_tensors.get(group_id, None)
      if unified_tensor is None:
        unified_tensor = UnifiedTensor(group_id, self.dtype)
      unified_tensor.append_cpu_tensor(cpu_part)
      self._unified_tensors[group_id] = unified_tensor