Extend type checking to all float datatypes

cebra/data/datasets.py

@@ -76,23 +76,35 @@ def __init__(
         device: str = "cpu"
     ):
         super().__init__(device=device)
-        self.neural = self._to_tensor(neural, torch.FloatTensor).float()
-        self.continuous = self._to_tensor(continuous, torch.FloatTensor)
-        self.discrete = self._to_tensor(discrete, torch.LongTensor)
+        self.neural = self._to_tensor(neural, check_dtype="float").float()
+        self.continuous = self._to_tensor(continuous, check_dtype="float").float()
+        self.discrete = self._to_tensor(discrete, check_dtype="integer")
         if self.continuous is None and self.discrete is None:
             raise ValueError(
                 "You have to pass at least one of the arguments 'continuous' or 'discrete'."
             )
         self.offset = offset
 
-    def _to_tensor(self, array, check_dtype=None):
+    def _to_tensor(self, array, check_dtype: str = None):
+        """Convert :py:func:`numpy.array` to :py:class:`torch.Tensor` if necessary and check the dtype.
+
+        Args:
+            array: Array to check.
+            check_dtype (list, optional): If not `None`, list of dtypes to which the values in `array`
+                must belong to. Defaults to None.
+
+        Returns:
+            The `array` as a :py:class:`torch.Tensor`.
+        """
         if array is None:
             return None
         if isinstance(array, np.ndarray):
             array = torch.from_numpy(array)
         if check_dtype is not None:
-            if not isinstance(array, check_dtype):
-                raise TypeError(f"{type(array)} instead of {check_dtype}.")
+            if (check_dtype == "integer" and not cebra.helper._is_integer(array)
+                ) or (check_dtype == "float" and not cebra.helper._is_floating(array)
+                ) or (check_dtype == "float_integer" and not cebra.helper._is_floating_or_integer(array)):
+                raise TypeError(f"{array.dtype} instead of {check_dtype}.")
         return array
 
     @property

cebra/helper.py

@@ -99,7 +99,7 @@ def _is_integer(y: Union[npt.NDArray, torch.Tensor]) -> bool:
 
 
 def _is_floating(y: Union[npt.NDArray, torch.Tensor]) -> bool:
-    """Check if the values in ``y`` are :py:class:`int`.
+    """Check if the values in ``y`` are :py:class:`float`.
 
     Note:
         There is no ``torch`` method to check that the ``dtype`` of a :py:class:`torch.Tensor`
@@ -118,6 +118,19 @@ def _is_floating(y: Union[npt.NDArray, torch.Tensor]) -> bool:
                 y, torch.Tensor) and torch.is_floating_point(y))
 
 
+def _is_floating_or_integer(y: Union[npt.NDArray, torch.Tensor]) -> bool:
+    """Check if the values in ``y`` are :py:class:`int` or :py:class:`float`.
+
+    Args:
+        y: An array, either as a :py:func:`numpy.array` or a :py:class:`torch.Tensor`.
+
+    Returns:
+        ``True`` if ``y`` contains :py:class:`float` or :py:class:`int`.
+    """
+
+    return _is_floating(y) or _is_integer(y)
+
+
 def get_loader_options(dataset: "cebra.data.Dataset") -> List[str]:
     """Return all possible dataloaders for the given dataset.