Unverified Commit 47bd9620 authored by Vasilis Vryniotis's avatar Vasilis Vryniotis Committed by GitHub
Browse files

Adding logging calls for raft and vit (#5044)

* Adding logging calls for raft and vit

* Linter fix
parent 588e9b5e
...@@ -8,6 +8,7 @@ from torch.nn.modules.batchnorm import BatchNorm2d ...@@ -8,6 +8,7 @@ from torch.nn.modules.batchnorm import BatchNorm2d
from torch.nn.modules.instancenorm import InstanceNorm2d from torch.nn.modules.instancenorm import InstanceNorm2d
from torchvision.ops import ConvNormActivation from torchvision.ops import ConvNormActivation
from ...utils import _log_api_usage_once
from ._utils import grid_sample, make_coords_grid, upsample_flow from ._utils import grid_sample, make_coords_grid, upsample_flow
...@@ -432,6 +433,7 @@ class RAFT(nn.Module): ...@@ -432,6 +433,7 @@ class RAFT(nn.Module):
If ``None`` (default), the flow is upsampled using interpolation. If ``None`` (default), the flow is upsampled using interpolation.
""" """
super().__init__() super().__init__()
_log_api_usage_once(self)
self.feature_encoder = feature_encoder self.feature_encoder = feature_encoder
self.context_encoder = context_encoder self.context_encoder = context_encoder
......
...@@ -11,6 +11,7 @@ import torch ...@@ -11,6 +11,7 @@ import torch
import torch.nn as nn import torch.nn as nn
from torch import Tensor from torch import Tensor
from ...utils import _log_api_usage_once
from ._api import WeightsEnum from ._api import WeightsEnum
from ._utils import handle_legacy_interface from ._utils import handle_legacy_interface
...@@ -139,6 +140,7 @@ class VisionTransformer(nn.Module): ...@@ -139,6 +140,7 @@ class VisionTransformer(nn.Module):
norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6), norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6),
): ):
super().__init__() super().__init__()
_log_api_usage_once(self)
torch._assert(image_size % patch_size == 0, "Input shape indivisible by patch size!") torch._assert(image_size % patch_size == 0, "Input shape indivisible by patch size!")
self.image_size = image_size self.image_size = image_size
self.patch_size = patch_size self.patch_size = patch_size
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment