def METHOD_NAME(cfg, args, rank, iterations, shared_results):
    is_cuda = torch.cuda.is_available()
    if is_cuda:
        torch.cuda.set_device(rank)
    model, loss_fn, optimizer = setup_model_loss_criterion(cfg, args, rank, is_cuda)
    for _ in range(iterations):
        input = torch.randn(1, args.input_size)
        target = torch.empty(args.batch_size, dtype=torch.long).random_(args.nb_classes)
        if is_cuda:
            input = input.cuda()
            target = target.cuda()
        train_step(input, target, model, loss_fn, optimizer)
    results = []
    for param in model.parameters():
        if len(results) == 0:
            results = param.flatten().cpu().data
        else:
            results = torch.cat((results, param.flatten().cpu().data), 0)
    shared_results[rank] = results