Update oceanmodelling.py

JaxGaussianProcesses · Mar 22, 2024 · 4be7629 · 4be7629
1 parent 5aa7cc2
commit 4be7629
Showing 1 changed file with 7 additions and 9 deletions.
diff --git a/docs/examples/oceanmodelling.py b/docs/examples/oceanmodelling.py
@@ -195,16 +195,14 @@ def dataset_3d(pos, vel):
 
 
 # %%
-
-
-@dataclass
 class VelocityKernel(gpx.kernels.AbstractKernel):
-    kernel0: gpx.kernels.AbstractKernel = field(
-        default_factory=lambda: gpx.kernels.RBF(active_dims=[0, 1])
-    )
-    kernel1: gpx.kernels.AbstractKernel = field(
-        default_factory=lambda: gpx.kernels.RBF(active_dims=[0, 1])
-    )
+    def __init__(
+        self,
+        kernel0: gpx.kernels.AbstractKernel = gpx.kernels.RBF(active_dims=[0, 1]),
+        kernel1: gpx.kernels.AbstractKernel = gpx.kernels.RBF(active_dims=[0, 1]),
+    ):
+        self.kernel0 = kernel0
+        self.kernel1 = kernel1
 
     def __call__(
         self, X: Float[Array, "1 D"], Xp: Float[Array, "1 D"]