Torch.jit.trace and devices

I cannot reproduce the issue using the posted code and just get a NamedTensor warning:

UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at  ../c10/core/TensorImpl.h:1257.)
  return torch.max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)

The traced model is successfully stored in the current working directory.