Thanks Mazhar_Shaikh, I chose that way, as it was the simplest implementation.
Just to give another example, the torchvision.models.detection.Generalized_RCNN model uses the same pattern, and complements it with a further check for model.training at line 44.

  
      github.com
  
  
    pytorch/vision/blob/8635be94d1216f10fb8302da89233bd86445e449/torchvision/models/detection/generalized_rcnn.py#L31-L62 28

def forward(self, images, targets=None):
    """
    Arguments:
        images (list[Tensor]): images to be processed
        targets (list[Dict[Tensor]]): ground-truth boxes present in the image (optional)


    Returns:
        result (list[BoxList] or dict[Tensor]): the output from the model.
            During training, it returns a dict[Tensor] which contains the losses.
            During testing, it returns list[BoxList] contains additional fields
            like `scores`, `labels` and `mask` (for Mask R-CNN models).


    """
    if self.training and targets is None:
        raise ValueError("In training mode, targets should be passed")
    original_image_sizes = [img.shape[-2:] for img in images]
    images, targets = self.transform(images, targets)
    features = self.backbone(images.tensors)
    if isinstance(features, torch.Tensor):
        features = OrderedDict([(0, features)])


  This file has been truncated. show original