fix: fill missing values in empirical_mean with 0, make grud more rub…

…ust;
WenjieDu · Jun 2, 2024 · 92d2766 · 92d2766
1 parent 4077383
commit 92d2766
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/pypots/imputation/grud/data.py b/pypots/imputation/grud/data.py
@@ -70,6 +70,8 @@ def __init__(
  self.empirical_mean = torch.sum(missing_mask * X, dim=[0, 1]) / torch.sum(
  missing_mask, dim=[0, 1]
  )
+ # fill nan with 0, in case some features have no observations
+ self.empirical_mean = torch.nan_to_num(self.empirical_mean, 0)
 
  def _fetch_data_from_array(self, idx: int) -> Iterable:
  """Fetch data according to index.