Fix: Map Advection-Diffusion model checkpoint to CPU

.ipynb_checkpoints/model_io_adv_dif-checkpoint.py

@@ -0,0 +1,22 @@
+import torch
+from dataclasses import dataclass, asdict
+from data_adv_dif import IntervalSplit
+from config_adv_dif import Config
+
+
+def save_model(path, model, tau_interval_split, alpha_interval_split, config):
+    torch.save({
+        'model_state_dict': model.state_dict(),
+        'alpha_interval_split': asdict(alpha_interval_split),
+        'tau_interval_split': asdict(tau_interval_split),
+        'config': asdict(config),
+    }, path)
+
+
+def load_model(path, model):
+    checkpoint = torch.load(path, map_location='cpu')
+    model.load_state_dict(checkpoint['model_state_dict'])
+    alpha_interval_split = IntervalSplit(**checkpoint['alpha_interval_split'])
+    tau_interval_split = IntervalSplit(**checkpoint['tau_interval_split'])
+    config = Config(**checkpoint['config'])
+    return model, alpha_interval_split, tau_interval_split, config

.ipynb_checkpoints/model_io_burgers-checkpoint.py

@@ -0,0 +1,22 @@
+import torch
+from dataclasses import dataclass, asdict
+from data_burgers import IntervalSplit
+from config_burgers import Config
+
+
+def save_model(path, model, tau_interval_split, re_interval_split, config):
+    torch.save({
+        'model_state_dict': model.state_dict(),
+        're_interval_split': asdict(re_interval_split),
+        'tau_interval_split': asdict(tau_interval_split),
+        'config': asdict(config),
+    }, path)
+
+
+def load_model(path, model):
+    checkpoint = torch.load(path, map_location='cpu')
+    model.load_state_dict(checkpoint['model_state_dict'])
+    re_interval_split = IntervalSplit(**checkpoint['re_interval_split'])
+    tau_interval_split = IntervalSplit(**checkpoint['tau_interval_split'])
+    config = Config(**checkpoint['config'])
+    return model, re_interval_split, tau_interval_split, config

app.ipynb

@@ -10,7 +10,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "* Running on local URL:  http://127.0.0.1:7885\n",
+      "* Running on local URL:  http://127.0.0.1:7886\n",
       "\n",
       "To create a public link, set `share=True` in `launch()`.\n"
      ]
@@ -18,7 +18,7 @@
     {
      "data": {
       "text/html": [
-       "<div><iframe src=\"http://127.0.0.1:7885/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
+       "<div><iframe src=\"http://127.0.0.1:7886/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
       ],
       "text/plain": [
        "<IPython.core.display.HTML object>"

model_io_adv_dif.py

@@ -14,7 +14,7 @@ def save_model(path, model, tau_interval_split, alpha_interval_split, config):
 
 
 def load_model(path, model):
-    checkpoint = torch.load(path)
+    checkpoint = torch.load(path, map_location='cpu')
     model.load_state_dict(checkpoint['model_state_dict'])
     alpha_interval_split = IntervalSplit(**checkpoint['alpha_interval_split'])
     tau_interval_split = IntervalSplit(**checkpoint['tau_interval_split'])

model_io_burgers.py

@@ -14,7 +14,7 @@ def save_model(path, model, tau_interval_split, re_interval_split, config):
 
 
 def load_model(path, model):
-    checkpoint = torch.load(path)
+    checkpoint = torch.load(path, map_location='cpu')
     model.load_state_dict(checkpoint['model_state_dict'])
     re_interval_split = IntervalSplit(**checkpoint['re_interval_split'])
     tau_interval_split = IntervalSplit(**checkpoint['tau_interval_split'])