Addressing code review comments.

Author: datumbox

references/classification/README.md

@@ -208,7 +208,7 @@ torchrun --nproc_per_node=8 train.py\
 ```
 
 Note that the above command corresponds to training on a single node with 8 GPUs.
-For generatring the pre-trained weights, we trained with 2 nodes, each with 8 GPUs (for a total of 32 GPUs),
+For generatring the pre-trained weights, we trained with 2 nodes, each with 8 GPUs (for a total of 16 GPUs),
 and `--batch_size 64`.
 
 ## Mixed precision training

torchvision/prototype/models/convnext.py

@@ -23,7 +23,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self.channels_last = kwargs.pop("channels_last", False)
         super().__init__(*args, **kwargs)
 
-    def forward(self, x):
+    def forward(self, x: Tensor) -> Tensor:
         #  TODO: Benchmark this against the approach described at https://github.com/pytorch/vision/pull/5197#discussion_r786251298
         if not self.channels_last:
             x = x.permute(0, 2, 3, 1)
@@ -34,7 +34,9 @@ def forward(self, x):
 
 
 class CNBlock(nn.Module):
-    def __init__(self, dim, layer_scale: float, stochastic_depth_prob: float, norm_layer: Callable[..., nn.Module]):
+    def __init__(
+        self, dim, layer_scale: float, stochastic_depth_prob: float, norm_layer: Callable[..., nn.Module]
+    ) -> None:
         super().__init__()
         self.block = nn.Sequential(
             ConvNormActivation(