Merge remote-tracking branch 'upstream/main' into batched-inference-and-padding

Author: CeliaBenquet

cebra/solver/base.py

@@ -37,8 +37,6 @@
 
 import literate_dataclasses as dataclasses
 import numpy.typing as npt
-import numpy as np
-import numpy.typing as npt
 import torch
 import torch.nn.functional as F
 from torch.utils.data import DataLoader
@@ -616,9 +614,9 @@ def transform(self,
         elif not isinstance(inputs, torch.Tensor):
             raise ValueError(
                 f"Inputs should be a torch.Tensor, not {type(inputs)}.")
-            
+
         self._check_is_fitted()
-        
+
         model, offset = self._select_model(inputs, session_id)
 
         if len(offset) < 2 and pad_before_transform:

tests/test_datasets.py

@@ -68,7 +68,6 @@ def test_demo():
 
 @pytest.mark.requires_dataset
 def test_hippocampus():
-
     pytest.skip("Outdated")
     dataset = cebra.datasets.init("rat-hippocampus-single")
     loader = cebra.data.ContinuousDataLoader(

tests/test_solver.py

@@ -195,7 +195,6 @@ def test_single_session(data_name, loader_initfunc, model_architecture,
         solver.load(temp_dir)
     _assert_equal(fitted_solver, solver)
 
-
     embedding = solver.transform(X)
     assert isinstance(embedding, torch.Tensor)
     assert embedding.shape == (X.shape[0], OUTPUT_DIMENSION)
@@ -224,9 +223,11 @@ def test_single_session(data_name, loader_initfunc, model_architecture,
     _assert_equal(fitted_solver, solver)
 
 
-@pytest.mark.parametrize("data_name, loader_initfunc, model_architecture, solver_initfunc",
-                         single_session_tests)
-def test_single_session_auxvar(data_name, loader_initfunc, model_architecture, solver_initfunc):
+@pytest.mark.parametrize(
+    "data_name, loader_initfunc, model_architecture, solver_initfunc",
+    single_session_tests)
+def test_single_session_auxvar(data_name, loader_initfunc, model_architecture,
+                               solver_initfunc):
 
     pytest.skip("Not yet supported")