Pytorch model split across different gpus

multi_gpu_model.py

from torch import nn

class Network(nn.Module):
    def __init__(self, split_gpus=False):
        super().__init__()
        self.module1 = ...
        self.module2 = ...
        self.split_gpus = split_gpus
        if split_gpus:            #considering only two gpus
            self.module1.cuda(0)
            self.module2.cuda(1)

    def forward(self, x):
        if self.split_gpus:
            x = x.cuda(0)
        x = self.module1(x)
        if self.split_gpus:
            x = x.cuda(1)
        x = self.module2(x)
        return x