Jyothirmai commited on
Commit
e0b6401
1 Parent(s): 931c795

Update models.py

Browse files
Files changed (1) hide show
  1. models.py +1 -1
models.py CHANGED
@@ -50,7 +50,7 @@ class VisualFeatureExtractor(nn.Module):
50
  vit_model = ViTModel(config)
51
 
52
  # Load the state dict specifically, excluding 'classifier.bias', 'classifier.weight'
53
- state_dict = torch.load('pytorch_model.bin', map_location=torch.device('cpu'))
54
  state_dict = {k: v for k, v in state_dict.items() if not k.startswith('classifier')}
55
  vit_model.load_state_dict(state_dict, strict=False)
56
 
 
50
  vit_model = ViTModel(config)
51
 
52
  # Load the state dict specifically, excluding 'classifier.bias', 'classifier.weight'
53
+ state_dict = torch.load('model/pytorch_model.bin', map_location=torch.device('cpu'))
54
  state_dict = {k: v for k, v in state_dict.items() if not k.startswith('classifier')}
55
  vit_model.load_state_dict(state_dict, strict=False)
56