diff --git a/Code/python/notebooks/predict.py b/Code/python/notebooks/predict.py
index 2947a53..0edc8ce 100644
--- a/Code/python/notebooks/predict.py
+++ b/Code/python/notebooks/predict.py
@@ -66,11 +66,9 @@ with torch.no_grad():
         y_pred, _ = model(X_t, X_f, X_c, warm // step, pred // step)
         lat_pred, lon_pred, alt_pred = denorm_coords(y_pred[...,0], y_pred[...,1], y_pred[...,2])
         lat_true, lon_true, alt_true = denorm_coords(y[...,0], y[...,1], y[...,2])
-        results.append((torch.stack([lat_true, lon_true, alt_true], dim=-1).cpu(), torch.stack([lat_pred, lon_pred, alt_pred], dim=-1).cpu()))
-    np_y_all = np.concatenate([
-        np.concatenate([t.numpy(), p.numpy()], axis=-1)
-        for t, p in results
-    ], axis=0)
+        results.append(torch.stack([lat_true, lon_true, alt_true, lat_pred, lon_pred, alt_pred], dim=1).cpu())
+    
+    np_y_all = torch.cat(results, dim=0).numpy()
     np.savetxt(
         f"predictions_{base_name}_wu{warm}_ps{pred}_aw{altitude_weight}.csv",
         np_y_all,