def setup_model_loss_criterion(cfg, args, rank, is_cuda):
    """
    setup model, criterion and optimizer based on input args
    """
    args.distributed_rank = rank
    cfg.distributed_training.distributed_rank = args.distributed_rank
    if cfg.distributed_training.distributed_world_size > 1:
        distributed_utils.distributed_init(cfg)
    torch.manual_seed(1)
    model = Model(args.input_size, args.nb_classes)
    loss_fn = nn.CrossEntropyLoss()
    if is_cuda:
        model = model.cuda()
        loss_fn = loss_fn.cuda()

    optimizer = optim.sgd.SGD(args, model.parameters())
    optimizer = optim.FairseqBMUF(cfg=cfg.bmuf, optimizer=optimizer)

    return model, loss_fn, optimizer
setup model, criterion and optimizer based on input args
tests/distributed/test_bmuf.py