Hi,
I find you could just implement one model class and use torch.nn.DataParallel to simply to train a model in parallel.
>>> net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])
>>> output = net(input_var)

https://pytorch.org/docs/stable/nn.html 82
for your reference.