I’m running the following code snippet on PyTorch version 1.6.0:
import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(3, 16, 1),
            nn.ReLU(),
        )

    def forward(self, model_inputs):
        return self.layers(model_inputs)


device = torch.device('cuda:0')
torch.cuda.set_device(device)

model = Model()
model = model.to(device=device, memory_format=torch.channels_last)

x = torch.zeros((1, 3, 32, 32), dtype=torch.float, device=device)
x = x.contiguous(memory_format=torch.channels_last)

loss = model(x).mean()
loss.backward()

During backward, it generates the following message:
[W TensorIterator.cpp:924] Warning: Mixed memory format inputs detected while calling the operator. The operator will output channels_last tensor even if some of the inputs are not in channels_last format. (function operator())

If I remove backward() call, then no warning is raised, as well as if I remove ReLU function from the model.
Where may formats not match? Can someone provide a deeper understanding of what’s happening here?