Adding _log_api_usage_once to Swin's reusable components. (#6175)

pytorch · Jun 16, 2022 · c7091ec · c7091ec
1 parent 1a4afa9
commit c7091ec
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/torchvision/models/swin_transformer.py b/torchvision/models/swin_transformer.py
@@ -34,6 +34,7 @@ class PatchMerging(nn.Module):
 
     def __init__(self, dim: int, norm_layer: Callable[..., nn.Module] = nn.LayerNorm):
         super().__init__()
+        _log_api_usage_once(self)
         self.dim = dim
         self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
         self.norm = norm_layer(4 * dim)
@@ -268,6 +269,7 @@ def __init__(
         attn_layer: Callable[..., nn.Module] = ShiftedWindowAttention,
     ):
         super().__init__()
+        _log_api_usage_once(self)
 
         self.norm1 = norm_layer(dim)
         self.attn = attn_layer(