add v in eval_burgers

alyashev · Feb 27, 2023 · 68db17b · 68db17b
1 parent d200910
commit 68db17b
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/train_burgers.py b/train_burgers.py
@@ -68,7 +68,7 @@ def test(config):
         ckpt = torch.load(ckpt_path)
         model.load_state_dict(ckpt['model'])
         print('Weights loaded from %s' % ckpt_path)
-    eval_burgers(model, dataloader, config, device)
+    eval_burgers(model, dataloader, dataset.v, config, device)
 
 
 if __name__ == '__main__':

diff --git a/train_utils/eval_2d.py b/train_utils/eval_2d.py
@@ -60,6 +60,7 @@ def eval_darcy(model,
 
 def eval_burgers(model,
                  dataloader,
+                 v,
                  config,
                  device,
                  use_tqdm=True):
@@ -78,7 +79,7 @@ def eval_burgers(model,
         out = model(x).reshape(y.shape)
         data_loss = myloss(out, y)
 
-        loss_u, f_loss = PINO_loss(out, x[:, 0, :, 0])
+        loss_u, f_loss = PINO_loss(out, x[:, 0, :, 0], v)
         test_err.append(data_loss.item())
         f_err.append(f_loss.item())