Fix NaN bug in validation (#59)

* Change from checking for NaNs to checking non finite values (Inf, -Inf, NaN)

* Fix - Check for non finite

* Remove debugging log

* Remove comment for another PR

---------

Co-authored-by: Louis Poulain--Auzéau <louis.poulainauzeau@meteoswiss.ch>

mlpp_lib/datasets.py

@@ -449,9 +449,9 @@ def drop_nans(self, group_size: int = 1):
         x, y, w = self._get_copies()
 
         event_axes = [self.dims.index(dim) for dim in self.dims if dim != "s"]
-        mask = da.any(da.isnan(da.from_array(x, name="x")), axis=event_axes)
+        mask = da.any(~da.isfinite(da.from_array(x, name="x")), axis=event_axes)
         if y is not None:
-            mask = mask | da.any(da.isnan(da.from_array(y, name="y")), axis=event_axes)
+            mask = mask | da.any(~da.isfinite(da.from_array(y, name="y")), axis=event_axes)
         mask = (~mask).compute()
 
         # with grouped samples, nans have to be removed in blocks: