add argument to sklearn api

Author: gonlairo

cebra/integrations/sklearn/cebra.py

@@ -1201,16 +1201,18 @@ def fit(
     def transform(self,
                   X: Union[npt.NDArray, torch.Tensor],
                   pad_before_transform: bool = True,
+                  batch_size: Optional[int] = None,
                   session_id: Optional[int] = None) -> npt.NDArray:
         """Transform an input sequence and return the embedding.
 
         Args:
             X: A numpy array or torch tensor of size ``time x dimension``.
             pad_before_transform: If ``False``, no padding is applied to the input sequence.
-            and the output sequence will be smaller than the input sequence due to the
-            receptive field of the model. If the input sequence is ``n`` steps long,
-            and a model with receptive field ``m`` is used, the output sequence would
-            only be ``n-m+1`` steps long.
+                and the output sequence will be smaller than the input sequence due to the
+                receptive field of the model. If the input sequence is ``n`` steps long,
+                and a model with receptive field ``m`` is used, the output sequence would
+                only be ``n-m+1`` steps long.
+            batch_size:
             session_id: The session ID, an :py:class:`int` between 0 and :py:attr:`num_sessions` for
                 multisession, set to ``None`` for single session.
 
@@ -1233,10 +1235,15 @@ def transform(self,
         # Input validation
         X = sklearn_utils.check_input_array(X, min_samples=len(self.offset_))
         input_dtype = X.dtype
+        #print(type(X))
+        #print(X.dtype)
 
         with torch.no_grad():
             output = self.solver_.transform(
-                X, pad_before_transform=pad_before_transform)
+                inputs=X,
+                pad_before_transform=pad_before_transform,
+                session_id=session_id,
+                batch_size=batch_size)
 
         if input_dtype == "float64":
             return output.astype(input_dtype)