Thanks for your answer!
This basically means there is no easy way and I would need to modify a copy of the following code, right?

  
      github.com
  
  
    pytorch/vision/blob/master/torchvision/models/segmentation/deeplabv3.py 9
import torch
from torch import nn
from torch.nn import functional as F

from ._utils import _SimpleSegmentationModel


__all__ = ["DeepLabV3"]


class DeepLabV3(_SimpleSegmentationModel):
    """
    Implements DeepLabV3 model from
    `"Rethinking Atrous Convolution for Semantic Image Segmentation"
    <https://arxiv.org/abs/1706.05587>`_.

    Arguments:
        backbone (nn.Module): the network used to compute the features for the model.
            The backbone should return an OrderedDict[Tensor], with the key being
            "out" for the last feature map used, and "aux" if an auxiliary classifier


  This file has been truncated. show original

  
  
    
    
  
  


This uses nn.sequential.
Do I also need to change this or does this “.to” work with nn.sequential (no separate forward function) as well?
Thanks!