Update model

model.py

@@ -3,6 +3,13 @@ import torchvision
 
 from torch import nn
 
+from torchvision.models._api import WeightsEnum
+from torch.hub import load_state_dict_from_url
+
+def get_state_dict(self, *args, **kwargs):
+    kwargs.pop("check_hash")
+    return load_state_dict_from_url(self.url, *args, **kwargs)
+WeightsEnum.get_state_dict = get_state_dict
 
 def create_model(num_classes:int=3, 
                         seed:int=42):