code: def demo_basic(rank, world_size):
      print(f"Running basic DDP example on rank {rank}.")
      setup(rank, world_size)
  
      if torch.cuda.is_available():
          device = rank
          device_ids = [rank]
      else:
          device = torch.device("cpu")
          device_ids = []
  
      # create model and move it to GPU with id rank
      model = ToyModel().to(device)
      ddp_model = DistributedDataParallel(model, device_ids=device_ids)
  
      with wandb.init(group="ddp-basic") as run:
          run.watch(models=ddp_model, log_freq=1, log_graph=True)
  
          loss_fn = nn.MSELoss()
          optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
  
          for _ in range(3):
              optimizer.zero_grad()
              outputs = ddp_model(torch.randn(20, 10))
              labels = torch.randn(20, 5).to(device)
              loss = loss_fn(outputs, labels)
              run.log({"loss": loss})
          loss.backward()
          optimizer.step()
  
      cleanup()
  
file_path: wandb/tests/functional_tests/t0_main/torch/t3_ddp_basic.py
imports: import os
  import torch
  import torch.distributed
  import torch.multiprocessing
  import torch.nn
  import torch.optim
  import wandb
  from torch.nn.parallel import DistributedDataParallel
label_desc: Use image node_id 3 for calling a global function with example usage: demo_basic(rank, world_size) without return types