need to implement forecast loss calculation in parallel

arnauqb · May 11, 2023 · 1d44bb5 · 1d44bb5
1 parent ee59e4c
commit 1d44bb5
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 5 deletions.
diff --git a/birds/forecast.py b/birds/forecast.py
@@ -65,6 +65,8 @@ def compute_forecast_loss(
         loss += loss_i
         n_samples_not_nan += 1
     if n_samples_not_nan == 0:
-        return torch.nan
-    return loss / n_samples_not_nan
+        loss = torch.nan
+    else:
+        loss = loss / n_samples_not_nan
+    return loss, loss # need to return it twice for the jacobian calculation
 
diff --git a/logo.png b/logo.png
diff --git a/test/test_forecast.py b/test/test_forecast.py
@@ -34,15 +34,15 @@ def test__compute_forecast_loss(self):
         observed_outputs = [torch.tensor(4.0)]
         assert compute_forecast_loss(
             loss_fn, model, parameter_generator, 5, observed_outputs
-        ) == 0
+        ) == (0,0)
         parameter_generator = lambda: torch.tensor(float("nan"))
         assert np.isnan(
             compute_forecast_loss(
                 loss_fn, model, parameter_generator, 5, observed_outputs
-            )
+            )[0]
         )
         parameter_generator = lambda: torch.tensor(2.0)
         model = lambda x: [x ** 3]
         assert compute_forecast_loss(
             loss_fn, model, parameter_generator, 5, observed_outputs
-        ) == (8-4)**2
+        ) == (16,16)