🐛 [torch.export] VGG model export.load() failure

[peri044]

Bug Description

With today's torch nightly, torch.export.load() on torchvision vgg 16 fails with the following error

torch._export.verifier.SpecViolationError: State dict entry for parameter features.0.weight is not an instance of torch.nn.Parameter.

To Reproduce

import torch
import torchvision.models as models 

model = models.vgg16(weights=None).eval().cuda()
random_inputs = [torch.rand([32, 3, 32, 32]).to("cuda")]
dim_x = torch.export.Dim("dim_x", min=1, max=32)
exp_program = torch.export.export(model, tuple(random_inputs), dynamic_shapes={"x" : {0 : dim_x}})
torch.export.save(exp_program, "./vgg_test.ep")
des_ep = torch.export.load("vgg_test.ep")
print(des_ep.module())

Expected behavior

Environment

Build information about Torch-TensorRT can be found by turning on debug messages

• Torch-TensorRT Version (e.g. 1.0.0):
• PyTorch Version (e.g. 1.0):
• CPU Architecture:
• OS (e.g., Linux):
• How you installed PyTorch (conda, pip, libtorch, source):
• Build command you used (if compiling from source):
• Are you using local sources or building from archives:
• Python version:
• CUDA version:
• GPU models and configuration:
• Any other relevant information:

Additional context

[peri044]

File on pytorch side

[peri044]

Fixed on pytorch side: pytorch/pytorch#113999