PyTorch conv2d doesn't propagate torch.channels_last memory format
When I try to use the torch.nn.conv2d operator on tensors which have the channels_last memory format, then the output does not keep this memory format.
I don't understand why, since conv2d is listed in the PyTorch wiki listing operators with channels last support.
Am I missing something?
Code to reproduce (tested with pytorch 1.6.0 and 1.7.0, on CPU, on Ubuntu 20.04):
import torch
import torch.nn.functional as F

N, C, H, W = 10, 4, 64, 128
out_channels = 2
kernel_size = (3, 3)

memory_format = torch.channels_last

tsr = torch.randn(N, C, H, W).to(memory_format=memory_format)
kernel = torch.randn(out_channels, C, *kernel_size).to(memory_format=memory_format)

conv_out = F.conv2d(tsr, kernel)

print(conv_out.is_contiguous(memory_format=memory_format)) # False


The conv2d operator is listed under the list of GPU operators supporting channels_last. This is not true for the CPU version of conv2d:
If you switch to cuda device, it will return True:
tsr = torch.randn(N, C, H, W).to('cuda', memory_format=memory_format)
kernel = torch.randn(out_channels, C, *kernel_size).to('cuda', memory_format=memory_format)
conv_out = F.conv2d(tsr, kernel)

&gt;&gt;&gt; conv_out.is_contiguous(memory_format=memory_format)
True

