Using a pre-trained PyTorch model (N C HW format) but my acceleration platform requires model in NHWC format.
Is there an easy way for converting PyTorch model to NHWC format?
I have permuted the weights by fetching weights from PyTorch’s state_dict() method like -

if('conv' in str(key)):
            params[key] = value.permute(0,2,3,1)

But am unable to repopulate the model with permuted dictionary - model.load_state_dict(Updated_params) gives -
size mismatch for stage4.2.branches.2.3.conv1.weight: copying a param with shape torch.Size([128, 3,
3, 128]) from checkpoint, the shape in current model is torch.Size([128, 128, 3, 3]).

To resolve this, how can I define layers in the new model in NHWC format in PyTorch.
The link below doesn’t seem to solve this issue.
Channels Last Memory Format in PyTorch 13
Thanks!