[docs] fix incorrect example in convert_conv3d_weight_memory_format (…

…#129318) The current example fails when using `torch.channels_last`, and the docs are slightly incorrect for the 3d case. Pull Request resolved: #129318 Approved by: https://github.com/albanD
pytorch · Aug 7, 2024 · 525fdc0 · 525fdc0
1 parent 6a348e5
commit 525fdc0
Showing 1 changed file with 8 additions and 8 deletions.
diff --git a/torch/nn/utils/memory_format.py b/torch/nn/utils/memory_format.py
@@ -16,13 +16,13 @@ def convert_conv2d_weight_memory_format(module, memory_format):
         layer with 4d weight will be affected by ``model.to``, which does not
         necessarily benefit from conversion to specified ``memory_format``.
         One place we are confident in is that NHWC(channels_last) conversion for
-        convolution in cuDNN, As it is beneficial to run convolution in NHWC,
+        convolution in cuDNN, as it is beneficial to run convolution in NHWC,
         even in cases where we have to apply permutation to input tensors.
 
         Hence our strategy here is to convert only the weight of convolution to
         channels_last. This ensures that;
         1. Fast convolution kernels will be used, the benefit of which could
-        outweigh overhead of permutation (if input is not in the same format)
+        outweigh overhead of permutation (if input is not in the same format).
         2. No unnecessary permutations are applied on layers that do not benefit
         from memory_format conversion.
 
@@ -85,18 +85,18 @@ def convert_conv3d_weight_memory_format(module, memory_format):
     provides considerable speed up for fp16 data on CUDA devices with compute capability >= 7.0
 
     .. note::
-        Calling ``model.to(memory_format=torch.channels_last)`` is more aggressive
+        Calling ``model.to(memory_format=torch.channels_last_3d)`` is more aggressive
         than the utility function ``convert_conv3d_weight_memory_format``. Any
         layer with 4d weight will be affected by ``model.to``, which does not
         necessarily benefit from conversion to specified ``memory_format``.
-        One place we are confident in is that NHWC(channels_last) conversion for
-        convolution in cuDNN, As it is beneficial to run convolution in NHWC,
+        One place we are confident in is that NDHWC(channels_last_3d) conversion for
+        convolution in cuDNN, as it is beneficial to run convolution in NDHWC,
         even in cases where we have to apply permutation to input tensors.
 
         Hence our strategy here is to convert only the weight of convolution to
-        channels_last. This ensures that;
+        channels_last_3d. This ensures that;
         1. Fast convolution kernels will be used, the benefit of which could
-        outweigh overhead of permutation (if input is not in the same format)
+        outweigh overhead of permutation (if input is not in the same format).
         2. No unnecessary permutations are applied on layers that do not benefit
         from memory_format conversion.
 
@@ -133,7 +133,7 @@ def convert_conv3d_weight_memory_format(module, memory_format):
         >>> model = nn.Sequential(
         >>>     nn.Conv3d(8, 4, 3)).cuda().half()
         >>> # This is identical to:
-        >>> # nn.utils.convert_conv3d_weight_memory_format(model, torch.channels_last)
+        >>> # nn.utils.convert_conv3d_weight_memory_format(model, torch.channels_last_3d)
         >>> model = nn.utils.convert_conv3d_weight_memory_format(model, torch.channels_last_3d)
         >>> out = model(input)
     """