Try like this:
from torch.utils import mkldnn as mkldnn_utils
# change input to be mkldnn layout
input_ = input.to_mkldnn()
# change weight to be mkldnn layout
model_ = mkldnn_utils.to_mkldnn(model)
# NB: output_ is mkldnn layout
output_ = model_(input_)
# change output_ back to default layout
output = output_.to_dense()