diff --git a/torchvision/models/segmentation/_utils.py b/torchvision/models/segmentation/_utils.py index 9780ef5b63c..0bbea5d3e81 100644 --- a/torchvision/models/segmentation/_utils.py +++ b/torchvision/models/segmentation/_utils.py @@ -5,6 +5,7 @@ from torch.nn import functional as F from ..._internally_replaced_utils import load_state_dict_from_url +from ...utils import _log_api_usage_once class _SimpleSegmentationModel(nn.Module): @@ -12,6 +13,7 @@ class _SimpleSegmentationModel(nn.Module): def __init__(self, backbone: nn.Module, classifier: nn.Module, aux_classifier: Optional[nn.Module] = None) -> None: super().__init__() + _log_api_usage_once(self) self.backbone = backbone self.classifier = classifier self.aux_classifier = aux_classifier diff --git a/torchvision/models/segmentation/lraspp.py b/torchvision/models/segmentation/lraspp.py index df4b21e2ee9..f6c2583cac1 100644 --- a/torchvision/models/segmentation/lraspp.py +++ b/torchvision/models/segmentation/lraspp.py @@ -4,6 +4,7 @@ from torch import nn, Tensor from torch.nn import functional as F +from ...utils import _log_api_usage_once from .. import mobilenetv3 from ..feature_extraction import create_feature_extractor from ._utils import _load_weights @@ -37,6 +38,7 @@ def __init__( self, backbone: nn.Module, low_channels: int, high_channels: int, num_classes: int, inter_channels: int = 128 ) -> None: super().__init__() + _log_api_usage_once(self) self.backbone = backbone self.classifier = LRASPPHead(low_channels, high_channels, num_classes, inter_channels)