Fix KernelFactory return types

Author: AdrianSosic

baybe/surrogates/gaussian_process/components/kernel.py

@@ -90,7 +90,7 @@ def _validate_parameter_kinds(self, parameters: Iterable[Parameter]) -> None:
     @override
     def __call__(
         self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
-    ) -> Kernel:
+    ) -> Kernel | GPyTorchKernel:
         """Construct the kernel, validating parameter kinds before construction."""
         if self.parameter_selector is not None:
             params = [p for p in searchspace.parameters if self.parameter_selector(p)]
@@ -103,7 +103,7 @@ def __call__(
     @abstractmethod
     def _make(
         self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
-    ) -> Kernel:
+    ) -> Kernel | GPyTorchKernel:
         """Construct the kernel."""
 
     def __attrs_post_init__(self):