Repository: microsoft/DeepSpeed
File: tests/unit/alexnet_model.py (Python)

Code before:
# Define this here because we cannot pickle local lambda functions
def cast_to_half(x):
    return x.half()


def cifar_trainset(fp16=False):
    torchvision = pytest.importorskip("torchvision", minversion="0.5.0")
    import torchvision.transforms as transforms

    transform_list = [
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ]
    if fp16:
        transform_list.append(torchvision.transforms.Lambda(cast_to_half))

    transform = transforms.Compose(transform_list)

    local_rank = get_accelerator().current_device()

    # Only one rank per machine downloads.
    dist.barrier()
    if local_rank != 0:
        dist.barrier()
    data_root = os.getenv("TEST_DATA_DIR", "/tmp/")
    download = False if os.getenv("CIFAR10_DATASET_PATH") else True
    trainset = torchvision.datasets.CIFAR10(root=os.getenv("CIFAR10_DATASET_PATH", default=os.path.join(data_root, "cifar10-data")),
                                            train=True,
                                            download=download,
                                            transform=transform)
    if local_rank == 0:
        dist.barrier()
    return trainset


def train_cifar(model, config, num_steps=400, average_dp_losses=True, fp16=True, seed=123):
    if required_torch_version(min_version=2.1):
        fork_kwargs = {"device_type": get_accelerator().device_name()}
    else:
        fork_kwargs = {}
    with get_accelerator().random().fork_rng(devices=[get_accelerator().current_device_name()], **fork_kwargs):
        ds_utils.set_random_seed(seed)

        # disable dropout
        model.eval()

        trainset = cifar_trainset(fp16=fp16)
        config['local_rank'] = dist.get_rank()

        with no_child_process_in_deepspeed_io():
            engine, _, _, _ = deepspeed.initialize(config=config,

Review comment (suggestion):
I recommend simplification like the following to improve readability. I have approved regardless. 
```suggestion
    if os.getenv("CIFAR10_DATASET_PATH"):
       data_root = os.getenv("CIFAR10_DATASET_PATH")
       download = False
    else:
       data_root = os.path.join(os.getenv("TEST_DATA_DIR", "/tmp"), "cifar10-data")
       download = True
```

Code after:
# Define this here because we cannot pickle local lambda functions
def cast_to_half(x):
    return x.half()


def cifar_trainset(fp16=False):
    torchvision = pytest.importorskip("torchvision", minversion="0.5.0")
    import torchvision.transforms as transforms

    transform_list = [
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ]
    if fp16:
        transform_list.append(torchvision.transforms.Lambda(cast_to_half))

    transform = transforms.Compose(transform_list)

    local_rank = get_accelerator().current_device()

    # Only one rank per machine downloads.
    dist.barrier()
    if local_rank != 0:
        dist.barrier()
    data_root = os.getenv("TEST_DATA_DIR", "/tmp/")
    if os.getenv("CIFAR10_DATASET_PATH"):
        data_root = os.getenv("CIFAR10_DATASET_PATH")
        download = False
    else:
        data_root = os.path.join(os.getenv("TEST_DATA_DIR", "/tmp"), "cifar10-data")
        download = True
    trainset = torchvision.datasets.CIFAR10(root=data_root, train=True, download=download, transform=transform)
    if local_rank == 0:
        dist.barrier()
    return trainset


def train_cifar(model, config, num_steps=400, average_dp_losses=True, fp16=True, seed=123):
    if required_torch_version(min_version=2.1):
        fork_kwargs = {"device_type": get_accelerator().device_name()}
    else:
        fork_kwargs = {}
    with get_accelerator().random().fork_rng(devices=[get_accelerator().current_device_name()], **fork_kwargs):
        ds_utils.set_random_seed(seed)

        # disable dropout
        model.eval()

        trainset = cifar_trainset(fp16=fp16)
        config['local_rank'] = dist.get_rank()