Use DenseNet121 with MKL-DNN

Hi there,

Consider the simple Python script as following:

import torch
import torchvision

from torch.utils import mkldnn as mkldnn_utils

net = torchvision.models.densenet121()
net = mkldnn_utils.to_mkldnn(net)

data = torch.randn(1, 3, 224, 224)
data = data.to_mkldnn()

_ = net(data)

We retrieve a DensNet 121 by torchvision and convert it to MKL-DNN data format. Besides, we generate a data and reorder it to MKL-DNN data layout.

When the last line of the script is executed, an exception is raised:

RuntimeError: Could not run 'aten::_cat' with arguments from the 'MkldnnCPUTensorId' backend. 'aten::_cat' is only available for these backends: [CPUTensorId, VariableTensorId].

Obviously, is not supported while backend is MKL-DNN.

Well, my problem is, how do I use DenseNet 121 (from torchvision) and also get the benefit of MKL-DNN?