fix(scalers): remove cast to float (#1115)

Nixtla · Aug 23, 2024 · dbf02a4 · dbf02a4
1 parent 6aed50e
commit dbf02a4
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/nbs/common.scalers.ipynb b/nbs/common.scalers.ipynb
@@ -128,7 +128,7 @@
     "    **Returns:**<br>\n",
     "    `x_median`: torch.Tensor with normalized values.\n",
     "    \"\"\"\n",
-    "    x_nan = x.float().masked_fill(mask<1, float(\"nan\"))\n",
+    "    x_nan = x.masked_fill(mask<1, float(\"nan\"))\n",
     "    x_median, _ = x_nan.nanmedian(dim=dim, keepdim=keepdim)\n",
     "    x_median = torch.nan_to_num(x_median, nan=0.0)\n",
     "    return x_median\n",
@@ -150,7 +150,7 @@
     "    **Returns:**<br>\n",
     "    `x_mean`: torch.Tensor with normalized values.\n",
     "    \"\"\"\n",
-    "    x_nan = x.float().masked_fill(mask<1, float(\"nan\"))\n",
+    "    x_nan = x.masked_fill(mask<1, float(\"nan\"))\n",
     "    x_mean = x_nan.nanmean(dim=dim, keepdim=keepdim)\n",
     "    x_mean = torch.nan_to_num(x_mean, nan=0.0)\n",
     "    return x_mean"

diff --git a/neuralforecast/common/_scalers.py b/neuralforecast/common/_scalers.py
@@ -26,7 +26,7 @@ def masked_median(x, mask, dim=-1, keepdim=True):
     **Returns:**<br>
     `x_median`: torch.Tensor with normalized values.
     """
-    x_nan = x.float().masked_fill(mask < 1, float("nan"))
+    x_nan = x.masked_fill(mask < 1, float("nan"))
     x_median, _ = x_nan.nanmedian(dim=dim, keepdim=keepdim)
     x_median = torch.nan_to_num(x_median, nan=0.0)
     return x_median
@@ -49,7 +49,7 @@ def masked_mean(x, mask, dim=-1, keepdim=True):
     **Returns:**<br>
     `x_mean`: torch.Tensor with normalized values.
     """
-    x_nan = x.float().masked_fill(mask < 1, float("nan"))
+    x_nan = x.masked_fill(mask < 1, float("nan"))
     x_mean = x_nan.nanmean(dim=dim, keepdim=keepdim)
     x_mean = torch.nan_to_num(x_mean, nan=0.0)
     return x_mean