Rick-29 commited on
Commit
170347f
1 Parent(s): 8bdc963

fix bug in loading effnet models

Browse files
Files changed (1) hide show
  1. model.py +9 -0
model.py CHANGED
@@ -1,4 +1,13 @@
1
 
 
 
 
 
 
 
 
 
 
2
  import torch
3
  import torchvision
4
 
 
1
 
2
+ from torchvision.models._api import WeightsEnum
3
+ from torch.hub import load_state_dict_from_url
4
+
5
+ def get_state_dict(self, *args, **kwargs):
6
+ kwargs.pop("check_hash")
7
+ return load_state_dict_from_url(self.url, *args, **kwargs)
8
+ WeightsEnum.get_state_dict = get_state_dict
9
+
10
+
11
  import torch
12
  import torchvision
13