Hey,
im trying to split ptsemseg’s frrn model to run on 2 GPU’s (Model Parallel) but in doing so i get the following error message.
NeuronalNetwork/env/lib/python3.6/site-packages/torch/nn/functional.py:2941: UserWarning: nn.functional.upsample is deprecated. Use nn.functional.interpolate instead.
  warnings.warn("nn.functional.upsample is deprecated. Use nn.functional.interpolate instead.")
Traceback (most recent call last):
  File "/NeuronalNetwork/train/train.py", line 239, in <module>
    train(cfg, writer, logger)
  File "/NeuronalNetwork/train/train.py", line 129, in train
    outputs = model(images.to('cuda:0'))
  File "/NeuronalNetwork/env/lib/python3.6/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/NeuronalNetwork/train/ptsemseg/models/MP_frrn.py", line 86, in forward
    y, z = getattr(self, key)(y_upsampled, z)
  File "/NeuronalNetwork/env/lib/python3.6/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/NeuronalNetwork/train/ptsemseg/models/utils.py", line 151, in forward
    y_prime = self.conv1(x)
  File "/NeuronalNetwork/env/lib/python3.6/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/NeuronalNetwork/train/ptsemseg/models/utils.py", line 75, in forward
    outputs = self.cbr_unit(inputs)
  File "/NeuronalNetwork/env/lib/python3.6/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/NeuronalNetwork/env/lib/python3.6/site-packages/torch/nn/modules/container.py", line 117, in forward
    input = module(input)
  File "/NeuronalNetwork/env/lib/python3.6/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/NeuronalNetwork/env/lib/python3.6/site-packages/torch/nn/modules/conv.py", line 419, in forward
    return self._conv_forward(input, self.weight)
  File "/NeuronalNetwork/env/lib/python3.6/site-packages/torch/nn/modules/conv.py", line 416, in _conv_forward
    self.padding, self.dilation, self.groups)
RuntimeError: Expected tensor for argument #1 'input' to have the same device as tensor for argument #2 'weight'; but device 1 does not equal 0 (while checking arguments for cudnn_convolution)

My class (mostly just copied from the existing model with some edit)
import torch

import torch.nn as nn

from ptsemseg.models.frrn import frrn

import torch.nn.functional as F

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")

        

        # 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

my model is on GPU 0
       >  model.train()


        labels = labels.to('cuda:0')

        # forward pass
        optimizer.zero_grad()
        outputs = model(images.to('cuda:0'))

        # backward pass
        labels = labels.to(outputs.device)
    
        loss = loss_fn(input=outputs, target=labels)


Does anyone have an idea what the problem might be?