Hi, currently I have the following train function:
def train(model, dataset, sampler, criterion, optimizer, scheduler, cfg):
    model = DataParallel(model.cuda())

    loader = DataLoader(dataset, bs=cfg.BS, num_workers=4, sampler=sampler)

    for epoch_idx in range(cfg.EPOCHS):
        for batch, targets in loader:
            preds = model(batch)
            loss = criterion(preds, targets)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            scheduler.step()

How can I convert this code to use DistributeDataParallel?
I looked in the tutorials, but they initialise model, dataset, etc. inside the train function. I can’t do that since I want the signature to remain the same and have the flexibility of defining model, dataset, etc. outside the train function.
Can I just pass all that using arguments in mp.spawn?