diff --git a/Code/python/notebooks/predict.py b/Code/python/notebooks/predict.py
index 1981e82..4f5adc8 100644
--- a/Code/python/notebooks/predict.py
+++ b/Code/python/notebooks/predict.py
@@ -72,9 +72,10 @@ with torch.no_grad():
         lat_true = lat_true.flatten()
         lon_true = lon_true.flatten()
         alt_true = alt_true.flatten()
-        results.append(torch.cat([lat_true, lon_true, alt_true, lat_pred, lon_pred, alt_pred], dim=-1).cpu())
+        results.append(torch.cat([lat_true, lon_true, alt_true, lat_pred, lon_pred, alt_pred], dim=1))
     
-    np_y_all = torch.cat(results, dim=0).numpy()
+    results_tensor = torch.cat(results, dim=0)
+    np_y_all = results_tensor.cpu().numpy()
     np.savetxt(
         f"predictions_{base_name}_wu{warm}_ps{pred}_aw{altitude_weight}.csv",
         np_y_all,