def traced_recv()

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

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

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

  return torch.distributed.untraced_recv(tensor, src, group, tag)