import torch
from torch.nn.parallel import DistributedDataParallel as DDP


def setup_model(model, device: torch.device):
    if device.type == "cuda":
        model = model.to(device)
        return DDP(model, device_ids=[device.index], output_device=device.index)

    model = model.to(device)
    return DDP(model)
