diff --git a/claymodel/finetune/embedder/factory.py b/claymodel/finetune/embedder/factory.py
index 529cbfbb..708b2141 100644
--- a/claymodel/finetune/embedder/factory.py
+++ b/claymodel/finetune/embedder/factory.py
@@ -107,8 +107,8 @@ def add_encodings(self, patches, time, latlon, gsd):
     def forward(self, datacube):
         cube, time, latlon, gsd, waves = (
             datacube["pixels"],  # [B C H W]
-            datacube["time"],  # [B 2]
-            datacube["latlon"],  # [B 2]
+            datacube["time"],  # [B 4]  # week_sin, week_cos, hour_sin, hour_cos
+            datacube["latlon"],  # [B 4]  # lat_sin, lat_cos, lon_sin, lon_cos
             datacube["gsd"],  # 1
             datacube["waves"],  # [N]
         )  # [B C H W]
diff --git a/claymodel/finetune/regression/factory.py b/claymodel/finetune/regression/factory.py
index 023eac71..05bc928d 100644
--- a/claymodel/finetune/regression/factory.py
+++ b/claymodel/finetune/regression/factory.py
@@ -100,8 +100,8 @@ def forward(self, datacube):
         """
         cube, time, latlon, gsd, waves = (
             datacube["pixels"],  # [B C H W]
-            datacube["time"],  # [B 2]
-            datacube["latlon"],  # [B 2]
+            datacube["time"],  # [B 4]  # week_sin, week_cos, hour_sin, hour_cos
+            datacube["latlon"],  # [B 4]  # lat_sin, lat_cos, lon_sin, lon_cos
             datacube["gsd"],  # 1
             datacube["waves"],  # [N]
         )
diff --git a/claymodel/finetune/segment/factory.py b/claymodel/finetune/segment/factory.py
index dcc07448..fee33a29 100644
--- a/claymodel/finetune/segment/factory.py
+++ b/claymodel/finetune/segment/factory.py
@@ -106,8 +106,8 @@ def forward(self, datacube):
         """
         cube, time, latlon, gsd, waves = (
             datacube["pixels"],  # [B C H W]
-            datacube["time"],  # [B 2]
-            datacube["latlon"],  # [B 2]
+            datacube["time"],  # [B 4]  # week_sin, week_cos, hour_sin, hour_cos
+            datacube["latlon"],  # [B 4]  # lat_sin, lat_cos, lon_sin, lon_cos
             datacube["gsd"],  # 1
             datacube["waves"],  # [N]
         )
diff --git a/claymodel/model.py b/claymodel/model.py
index 300eb823..50c22b55 100644
--- a/claymodel/model.py
+++ b/claymodel/model.py
@@ -161,8 +161,8 @@ def mask_out(self, patches):
     def forward(self, datacube):
         cube, time, latlon, gsd, waves = (
             datacube["pixels"],  # [B C H W]
-            datacube["time"],  # [B 2]
-            datacube["latlon"],  # [B 2]
+            datacube["time"],  # [B 4]  # week_sin, week_cos, hour_sin, hour_cos
+            datacube["latlon"],  # [B 4]  # lat_sin, lat_cos, lon_sin, lon_cos
             datacube["gsd"],  # 1
             datacube["waves"],  # [N]
         )  # [B C H W]
diff --git a/docs/references.bib b/docs/references.bib
index 2b5c34cd..e9cee75d 100644
--- a/docs/references.bib
+++ b/docs/references.bib
@@ -1,2 +1,2 @@
 % This file contains bibliography references for the Clay Foundation Model documentation
-% Currently empty but required by the Jupyter Book configuration
\ No newline at end of file
+% Currently empty but required by the Jupyter Book configuration
diff --git a/docs/tutorials/wall-to-wall.ipynb b/docs/tutorials/wall-to-wall.ipynb
index 0b79f64f..da527826 100644
--- a/docs/tutorials/wall-to-wall.ipynb
+++ b/docs/tutorials/wall-to-wall.ipynb
@@ -482,12 +482,12 @@
     "datacube = {\n",
     "    \"platform\": platform,\n",
     "    \"time\": torch.tensor(\n",
-    "        np.hstack((week_norm, hour_norm)),\n",
+    "        np.column_stack((week_norm, hour_norm)),\n",
     "        dtype=torch.float32,\n",
     "        device=device,\n",
     "    ),\n",
     "    \"latlon\": torch.tensor(\n",
-    "        np.hstack((lat_norm, lon_norm)), dtype=torch.float32, device=device\n",
+    "        np.column_stack((lat_norm, lon_norm)), dtype=torch.float32, device=device\n",
     "    ),\n",
     "    \"pixels\": pixels.to(device),\n",
     "    \"gsd\": torch.tensor(stack.gsd.values, device=device),\n",