remove float16

cebra/integrations/sklearn/cebra.py

@@ -1235,7 +1235,7 @@ def transform(self,
         # Input validation
         #TODO: if inputs are in cuda, then it throws an error, deal with this.
         X = sklearn_utils.check_input_array(X, min_samples=len(self.offset_))
-        input_dtype = X.dtype
+        #input_dtype = X.dtype
 
         if isinstance(X, np.ndarray):
             X = torch.from_numpy(X)
@@ -1248,10 +1248,11 @@ def transform(self,
                 session_id=session_id,
                 batch_size=batch_size)
 
-        if input_dtype == "float64":
-            return output.astype(input_dtype)
+        #TODO: check if this is safe.
+        return output.numpy(force=True)
 
-        return output
+        #if input_dtype == "float64":
+        #    return output.astype(input_dtype)
 
     def fit_transform(
         self,

cebra/integrations/sklearn/utils.py

@@ -78,7 +78,8 @@ def check_input_array(X: npt.NDArray, *, min_samples: int) -> npt.NDArray:
         X,
         accept_sparse=False,
         accept_large_sparse=False,
-        dtype=("float16", "float32", "float64"),
+        # NOTE: remove float16 because F.pad does not allow float16.
+        dtype=("float32", "float64"),
         order=None,
         copy=False,
         force_all_finite=True,