fix: turn missing_mask into torch.float;

WenjieDu · Sep 21, 2023 · 18945e8 · 18945e8
1 parent acfe255
commit 18945e8
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/pypots/data/base.py b/pypots/data/base.py
@@ -205,7 +205,7 @@ def _fetch_data_from_array(self, idx: int) -> Iterable:
         """
 
         X = self.X[idx].to(torch.float32)
-        missing_mask = ~torch.isnan(X)
+        missing_mask = (~torch.isnan(X)).to(torch.float32)
         X = torch.nan_to_num(X)
         sample = [
             torch.tensor(idx),
@@ -280,7 +280,7 @@ def _fetch_data_from_file(self, idx: int) -> Iterable:
             self.file_handle = self._open_file_handle()
 
         X = torch.from_numpy(self.file_handle["X"][idx]).to(torch.float32)
-        missing_mask = ~torch.isnan(X)
+        missing_mask = (~torch.isnan(X)).to(torch.float32)
         X = torch.nan_to_num(X)
         sample = [
             torch.tensor(idx),

diff --git a/pypots/imputation/gpvae/data.py b/pypots/imputation/gpvae/data.py
@@ -116,7 +116,7 @@ def _fetch_data_from_file(self, idx: int) -> Iterable:
             self.file_handle = self._open_file_handle()
 
         X = torch.from_numpy(self.file_handle["X"][idx]).to(torch.float32)
-        missing_mask = ~torch.isnan(X)
+        missing_mask = (~torch.isnan(X)).to(torch.float32)
         X = torch.nan_to_num(X)
 
         sample = [

diff --git a/tests/imputation/gpvae.py b/tests/imputation/gpvae.py
@@ -80,7 +80,7 @@ def test_2_parameters(self):
             and self.gp_vae.best_model_dict is not None
         )
 
-    @pytest.mark.xdist_group(name="imputation-GPVAE")
+    @pytest.mark.xdist_group(name="imputation-gpvae")
     def test_3_saving_path(self):
         # whether the root saving dir exists, which should be created by save_log_into_tb_file
         assert os.path.exists(