import torch


def data_parallel_wrapper(model, cur_device, cfg):
    """Wraps a given model into pytorch data parallel.

    Args:
        model (torch.nn.Module): Image classifier.
        cur_device (int): gpu_id
        cfg : Reference to config.

    Returns:
        torch.nn.Dataparallel: model wrapped in dp.
    """
    model.cuda(cur_device)
    # model = torch.nn.DataParallel(model, device_ids = [cur_device])
    # assert cfg.NUM_GPUS == torch.cuda.device_count(), f"Expected torch device count (i.e {torch.cuda.device_count()}) same as number of gpus (i.e {cfg.NUM_GPUS}) in config file"
    model = torch.nn.DataParallel(
        model, device_ids=[i for i in range(torch.cuda.device_count())]
    )
    return model