1import trace_layer2 as models
2import torch
3
4x=torch.randn(1, 3, 224, 224)
5
6state_dict = torch.load('swav_imagenet_layer2.pt', map_location='cpu')
7
8model = models.resnet50w2()
9model.load_state_dict(state_dict)
10model.eval()
11feature = model(x)
12
13traced_model = torch.jit.load('traced_swav_imagenet_layer2.pt', map_location='cpu')
14traced_model.eval()
15feature = traced_model(x)
16