fix cov-update in predict

strongio · Jan 24, 2025 · eb77787 · eb77787
1 parent d3f8d0d
commit eb77787
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/torchcast/exp_smooth/exp_smooth.py b/torchcast/exp_smooth/exp_smooth.py
@@ -57,7 +57,11 @@ def predict(self,
         new_mean = update_tensor(mean, new=(F @ mean[mask].unsqueeze(-1)).squeeze(-1), mask=mask)
         new_cov = kwargs['cov1step']
         if len(cov.shape):  # see note in _update() above
-            new_cov = update_tensor(new_cov.clone(), new=F @ cov[mask] @ F.permute(0, 2, 1), mask=mask)
+            new_cov = update_tensor(
+                orig=new_cov,
+                new=new_cov[mask] + F @ cov[mask] @ F.permute(0, 2, 1),
+                mask=mask
+            )
 
         return new_mean, new_cov