```python import trace_layer2 as models import torch x=torch.randn(1, 3, 224, 224) state_dict = torch.load('swav_imagenet_layer2.pt', map_location='cpu') model = models.resnet50w2() model.load_state_dict(state_dict) model.eval() feature = model(x) traced_model = torch.jit.load('traced_swav_imagenet_layer2.pt', map_location='cpu') traced_model.eval() feature = traced_model(x) ```