moved GriddedDataset export to from_numpy and tests pass.

MPoL-dev · Dec 27, 2023 · 386108c · 386108c
1 parent f8e4018
commit 386108c
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/src/mpol/gridding.py b/src/mpol/gridding.py
@@ -12,6 +12,8 @@
 import numpy.typing as npt
 from fast_histogram import histogram as fast_hist
 
+import torch
+
 from mpol.coordinates import GridCoords
 from mpol.exceptions import DataError, ThresholdExceededError, WrongDimensionError
 from mpol.datasets import GriddedDataset
@@ -657,9 +659,9 @@ def to_pytorch_dataset(
         return GriddedDataset(
             coords=self.coords,
             nchan=self.nchan,
-            vis_gridded=self.vis_gridded,
-            weight_gridded=self.weight_gridded,
-            mask=self.mask,
+            vis_gridded=torch.from_numpy(self.vis_gridded),
+            weight_gridded=torch.from_numpy(self.weight_gridded),
+            mask=torch.from_numpy(self.mask),
         )