took some time to boil it down:
import torch

import torch.optim as optim

import torch.nn.functional as F

from torch.utils import data

from ptsemseg.models.frrn import frrn

from ptsemseg.loss import cross_entropy2d

class ModelParallelfrrn(frrn):

    def __init__(self, *args, **kwargs):

        super(ModelParallelfrrn, self).__init__(

            n_classes=21, model_type="B", group_norm=False, n_groups=16)  # default init values

        self.conv1.to("cuda:0")

        self.up_residual_units.to("cuda:0")

        self.split_conv.to("cuda:0")

        self.merge_conv.to("cuda:1")

        self.down_residual_units.to("cuda:1")

        self.classif_conv.to("cuda:1")

    def forward(self, x):

        # pass to initial conv

        x = self.conv1(x.to("cuda:0"))

        # pass through residual units

        for i in range(3):

            x = self.up_residual_units[i](x)

        # divide stream

        y = x                               # full image resolution stream

        z = self.split_conv(x) # processed image stream

        prev_channels = 48

        # encoding

        for n_blocks, channels, scale in self.encoder_frru_specs:

            # maxpool bigger feature map

            y_pooled = F.max_pool2d(y, stride=2, kernel_size=2, padding=0)

            # pass through encoding FRRUs

            for block in range(n_blocks):

                key = "_".join(map(str, ["encoding_frru", n_blocks, channels, scale, block]))

                y, z = getattr(self, key)(y_pooled, z)

            prev_channels = channels

        # move both streams to GPU 1

        y = y.to("cuda:1")

        z = z.to("cuda:1")

        print("\nINside\n")

        for name, param in  self.named_parameters():

                print(name, param.device)

        # decoding

        for n_blocks, channels, scale in self.decoder_frru_specs:

            # bilinear upsample smaller feature map

            upsample_size = torch.Size([_s * 2 for _s in y.size()[-2:]])

            y_upsampled = F.upsample(y, size=upsample_size, mode="bilinear", align_corners=True).to("cuda:1")

            # pass through decoding FRRUs

            for block in range(n_blocks):

                key = "_".join(map(str, ["decoding_frru", n_blocks, channels, scale, block]))

                # print("Incoming FRRU Size: ", key, y_upsampled.shape, z.shape)

                y, z = getattr(self, key)(y_upsampled, z)

                # print("Outgoing FRRU Size: ", key, y.shape, z.shape)

            prev_channels = channels

        # merge streams

        x = torch.cat(

            [F.upsample(y, scale_factor=2, mode="bilinear", align_corners=True), z], dim=1

        ).to("cuda:1")

        x = self.merge_conv(x)

        # pass through residual units

        for i in range(3):

            x = self.down_residual_units[i](x)

        # final 1x1 conv to get classification

        x = self.classif_conv(x)

        return x

def train():

    print("Available GPU's:{}".format(torch.cuda.device_count()))

    num_classes = 1000

    num_batches = 3

    batch_size = 20

    image_w = 128

    image_h = 128

    # Setup Model

    model = ModelParallelfrrn().to('cuda:0')

    for name, param in  model.named_parameters():

            print(name, param.device)

    optimizer = optim.SGD(model.parameters(), 1.0e-10,weight_decay= 0.0005, momentum=0.99)

    one_hot_indices = torch.LongTensor(batch_size) \

                            .random_(0, num_classes) \

                            .view(batch_size, 1)

    for _ in range(num_batches):

    

    # generate random inputs and labels

        inputs = torch.randn(batch_size, 3, image_w, image_h)

        labels = torch.zeros(batch_size, num_classes) \

                        .scatter_(1, one_hot_indices, 1)

        # run forward pass

        optimizer.zero_grad()

        outputs = model(inputs.to('cuda:0'))

        # run backward pass

        labels = labels.to(outputs.device)

        loss_fn = cross_entropy2d(outputs, labels, size_average=False).backward()

        optimizer.step()

if __name__ == "__main__":

    train()