def traced_send()

in sample_workloads/lit-gpt-demo/utilities/monitor_collectives.py [0:0]

• 5 lines of code
• 2 McCabe index (conditional complexity)

def traced_send(tensor, dst, group=None, tag=0):
  """Intercepts invocations of torch.distributed.send.
  """
  if _should_rank_record_comm(group, peer_rank=dst, is_ring=False):
    message_size = tensor.nelement() * tensor.element_size()
    _emit_call_description('send', message_size, group, dst)

  return torch.distributed.untraced_send(tensor, dst, group, tag)