allow loading checkpoint from url

microsoft · Feb 15, 2023 · 99d0b4d · 99d0b4d
1 parent d6dec5f
commit 99d0b4d
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/src/climax/global_forecast/module.py b/src/climax/global_forecast/module.py
@@ -55,8 +55,10 @@ def __init__(
             self.load_pretrained_weights(pretrained_path)
 
     def load_pretrained_weights(self, pretrained_path):
-        checkpoint = torch.load(pretrained_path, map_location=torch.device("cpu"))
-
+        if pretrained_path.startswith('http'):
+            checkpoint = torch.hub.load_state_dict_from_url(pretrained_path)
+        else:
+            checkpoint = torch.load(pretrained_path, map_location=torch.device("cpu"))
         print("Loading pre-trained checkpoint from: %s" % pretrained_path)
         checkpoint_model = checkpoint["state_dict"]
         # interpolate positional embedding

diff --git a/src/climax/regional_forecast/module.py b/src/climax/regional_forecast/module.py
@@ -55,7 +55,10 @@ def __init__(
             self.load_pretrained_weights(pretrained_path)
 
     def load_pretrained_weights(self, pretrained_path):
-        checkpoint = torch.load(pretrained_path, map_location=torch.device("cpu"))
+        if pretrained_path.startswith('http'):
+            checkpoint = torch.hub.load_state_dict_from_url(pretrained_path)
+        else:
+            checkpoint = torch.load(pretrained_path, map_location=torch.device("cpu"))
 
         print("Loading pre-trained checkpoint from: %s" % pretrained_path)
         checkpoint_model = checkpoint["state_dict"]