AFAIK, самый простой способ выполнить распределенное обучение (несколько режимов) с помощью Pytorch - это что-то вроде:
sampler = torch.utils.data.distributed.DistributedSampler(train_data)
data_loader = torch.utils.data.DataLoader(dataset, sampler=sampler)
model = torch.nn.DataParallel(model).cuda()
for data, target in data_loader:
out = model(data)
...
Но что, если у меня уже есть большие тензорные данные в руках, и я хочу разделить ираспространять его и получать тот же вывод, что и приведенный выше фрагмент? В частности,
model = torch.nn.DataParallel(model).cuda()
data = do_sth_fuct(data)
out = model(data)
Существует ли PyTorch API для этого? Иначе, как лучше всего это сделать? Заранее спасибо!