diff --git a/bluemath_tk/deeplearning/_base_model.py b/bluemath_tk/deeplearning/_base_model.py
index 1aafe24..c764f4e 100644
--- a/bluemath_tk/deeplearning/_base_model.py
+++ b/bluemath_tk/deeplearning/_base_model.py
@@ -1,3 +1,4 @@
+import copy
 from abc import abstractmethod
 from typing import Dict, Optional, Union
 
@@ -202,7 +203,7 @@ def fit(
             if val_loss < best_val_loss:
                 best_val_loss = val_loss
                 patience_counter = 0
-                best_model_state = self.model.state_dict().copy()
+                best_model_state = copy.deepcopy(self.model.state_dict())
             else:
                 patience_counter += 1
                 if patience_counter >= patience:
diff --git a/bluemath_tk/deeplearning/autoencoders.py b/bluemath_tk/deeplearning/autoencoders.py
index ad49bb9..6583db8 100644
--- a/bluemath_tk/deeplearning/autoencoders.py
+++ b/bluemath_tk/deeplearning/autoencoders.py
@@ -20,6 +20,7 @@
 - evaluate(X)
 """
 
+import copy
 from typing import Dict, Optional, Tuple
 
 import numpy as np
@@ -91,17 +92,20 @@ def __init__(
 
     def _build_model(self, input_shape: Tuple, **kwargs) -> nn.Module:
         """Build the standard fully-connected autoencoder model."""
-        # Handle input shape: (n_samples, n_features) or (n_features,)
+        # Handle input shape: (n_samples, ...) or a single sample shape.
+        # For fit(X), input_shape includes the batch dimension, so the
+        # per-sample shape is input_shape[1:].
         if len(input_shape) == 1:
-            n_features = input_shape[0]
+            sample_shape = (input_shape[0],)
         else:
-            # Take last dimension as features (handles (n_samples, n_features))
-            n_features = input_shape[-1]
+            sample_shape = tuple(input_shape[1:])
+        n_features = int(np.prod(sample_shape))
 
         class StandardAutoencoderModel(nn.Module):
-            def __init__(self, n_features, hidden_dims, k):
+            def __init__(self, n_features, hidden_dims, k, sample_shape):
                 super().__init__()
                 self.n_features = n_features
+                self.sample_shape = tuple(sample_shape)
                 # Encoder
                 encoder_layers = []
                 prev_dim = n_features
@@ -132,7 +136,7 @@ def forward(self, x):
                     x = x.unsqueeze(0)
                 z = self.encoder(x)
                 x_recon = self.decoder(z)
-                return x_recon
+                return x_recon.view(x_recon.size(0), *self.sample_shape)
 
             def encode_forward(self, x):
                 """Encode input to latent space."""
@@ -142,7 +146,7 @@ def encode_forward(self, x):
                     x = x.unsqueeze(0)
                 return self.encoder(x)
 
-        return StandardAutoencoderModel(n_features, self.hidden_dims, self.k)
+        return StandardAutoencoderModel(n_features, self.hidden_dims, self.k, sample_shape)
 
 
 class OrthogonalAutoencoder(BaseDeepLearningModel):
@@ -204,16 +208,22 @@ def __init__(
 
     def _build_model(self, input_shape: Tuple, **kwargs) -> nn.Module:
         """Build the orthogonal autoencoder model."""
-        # Handle input shape: (n_samples, n_features) or (n_features,)
+        # Handle input shape: (n_samples, ...) or a single sample shape.
+        # For fit(X), input_shape includes the batch dimension, so the
+        # per-sample shape is input_shape[1:].
         if len(input_shape) == 1:
-            n_features = input_shape[0]
+            sample_shape = (input_shape[0],)
         else:
-            n_features = input_shape[-1]
+            sample_shape = tuple(input_shape[1:])
+        n_features = int(np.prod(sample_shape))
 
         class OrthogonalAutoencoderModel(nn.Module):
-            def __init__(self, n_features, hidden_dims, k, lambda_W, lambda_Z):
+            def __init__(
+                self, n_features, hidden_dims, k, lambda_W, lambda_Z, sample_shape
+            ):
                 super().__init__()
                 self.n_features = n_features
+                self.sample_shape = tuple(sample_shape)
                 self.lambda_W = lambda_W
                 self.lambda_Z = lambda_Z
 
@@ -265,7 +275,7 @@ def forward(self, x):
                 z = z + 0 * ortho_loss
 
                 x_recon = self.decoder(z)
-                return x_recon
+                return x_recon.view(x_recon.size(0), *self.sample_shape)
 
             def encode_forward(self, x):
                 """Encode input to latent space."""
@@ -287,7 +297,12 @@ def get_regularization_losses(self):
                 return ortho_loss, decorr_loss
 
         return OrthogonalAutoencoderModel(
-            n_features, self.hidden_dims, self.k, self.lambda_W, self.lambda_Z
+            n_features,
+            self.hidden_dims,
+            self.k,
+            self.lambda_W,
+            self.lambda_Z,
+            sample_shape,
         )
 
     def fit(
@@ -410,7 +425,7 @@ def fit(
             if val_loss < best_val_loss:
                 best_val_loss = val_loss
                 patience_counter = 0
-                best_model_state = self.model.state_dict().copy()
+                best_model_state = copy.deepcopy(self.model.state_dict())
             else:
                 patience_counter += 1
                 if patience_counter >= patience:
@@ -838,6 +853,7 @@ def __init__(
                 self.H = H
                 self.W = W
                 self.C = C
+                self.d_model = d_model
 
                 # Patchify + embed + pos
                 self.patchify = Patchify(patch_size)
@@ -879,6 +895,7 @@ def __init__(
                         )
                     )
                 self.decoder_blocks = nn.Sequential(*decoder_blocks)
+                self.patch_reconstruct = nn.Linear(d_model, Pdim)
 
                 # Reconstruct patches
                 self.unpatchify = Unpatchify(patch_size, Hp, Wp, C)
@@ -930,7 +947,8 @@ def forward(self, x):
                     y = block(y)
 
                 # Reconstruct patches
-                rec_patches = self.unpatchify(y)  # (B, C, H+pad_h, W+pad_w)
+                patch_tokens = self.patch_reconstruct(y)  # (B, N, Pdim)
+                rec_patches = self.unpatchify(patch_tokens)  # (B, C, H+pad_h, W+pad_w)
 
                 # Crop padding
                 if self.pad_h > 0 or self.pad_w > 0:
@@ -1036,6 +1054,47 @@ def __init__(
         self.k = k
         super().__init__(device=device, **kwargs)
 
+    def fit(
+        self,
+        X: np.ndarray,
+        y: Optional[np.ndarray] = None,
+        validation_split: float = 0.2,
+        epochs: int = 500,
+        batch_size: int = 64,
+        learning_rate: float = 1e-3,
+        optimizer: Optional[torch.optim.Optimizer] = None,
+        criterion: Optional[nn.Module] = None,
+        patience: int = 20,
+        verbose: int = 1,
+        **kwargs,
+    ) -> Dict[str, list]:
+        """Fit the ConvLSTM autoencoder.
+
+        If ``y`` is not provided, the model reconstructs the last frame of each
+        input sequence, matching the documented prediction shape ``(B, C, H, W)``.
+        """
+        if y is None:
+            if X.ndim != 5:
+                raise ValueError(
+                    "ConvLSTMAutoencoder expects 5D input "
+                    "(n_samples, seq_len, C, H, W)."
+                )
+            y = X[:, -1]
+
+        return super().fit(
+            X,
+            y=y,
+            validation_split=validation_split,
+            epochs=epochs,
+            batch_size=batch_size,
+            learning_rate=learning_rate,
+            optimizer=optimizer,
+            criterion=criterion,
+            patience=patience,
+            verbose=verbose,
+            **kwargs,
+        )
+
     def _build_model(self, input_shape: Tuple, **kwargs) -> nn.Module:
         """Build the ConvLSTM autoencoder model."""
         # Parse input shape: (n_samples, seq_len, C, H, W) - channels-first format
@@ -1068,7 +1127,7 @@ def __init__(self, seq_len, H, W, C, k, pad_h, pad_w):
                 self.convlstm1 = ConvLSTM(
                     input_dim=C,
                     hidden_dim=32,
-                    kernel_size=3,
+                    kernel_size=(3, 3),
                     num_layers=1,
                     batch_first=True,
                     return_all_layers=False,
@@ -1077,7 +1136,7 @@ def __init__(self, seq_len, H, W, C, k, pad_h, pad_w):
                 self.convlstm2 = ConvLSTM(
                     input_dim=32,
                     hidden_dim=32,
-                    kernel_size=3,
+                    kernel_size=(3, 3),
                     num_layers=1,
                     batch_first=True,
                     return_all_layers=False,
@@ -1276,6 +1335,47 @@ def __init__(
         self.k = k
         super().__init__(device=device, **kwargs)
 
+    def fit(
+        self,
+        X: np.ndarray,
+        y: Optional[np.ndarray] = None,
+        validation_split: float = 0.2,
+        epochs: int = 500,
+        batch_size: int = 64,
+        learning_rate: float = 1e-3,
+        optimizer: Optional[torch.optim.Optimizer] = None,
+        criterion: Optional[nn.Module] = None,
+        patience: int = 20,
+        verbose: int = 1,
+        **kwargs,
+    ) -> Dict[str, list]:
+        """Fit the hybrid ConvLSTM-Transformer autoencoder.
+
+        If ``y`` is not provided, the model reconstructs the last frame of each
+        input sequence, matching the documented prediction shape ``(B, C, H, W)``.
+        """
+        if y is None:
+            if X.ndim != 5:
+                raise ValueError(
+                    "HybridConvLSTMTransformerAutoencoder expects 5D input "
+                    "(n_samples, seq_len, C, H, W)."
+                )
+            y = X[:, -1]
+
+        return super().fit(
+            X,
+            y=y,
+            validation_split=validation_split,
+            epochs=epochs,
+            batch_size=batch_size,
+            learning_rate=learning_rate,
+            optimizer=optimizer,
+            criterion=criterion,
+            patience=patience,
+            verbose=verbose,
+            **kwargs,
+        )
+
     def _build_model(self, input_shape: Tuple, **kwargs) -> nn.Module:
         """Build the hybrid autoencoder model."""
         # Parse input shape: (n_samples, seq_len, C, H, W) - channels-first format
@@ -1314,6 +1414,7 @@ def __init__(
                 self.H = H
                 self.W = W
                 self.C = C
+                self.efficient_attention = efficient_attention
 
                 # ConvLSTM stack
                 from .layers import ConvLSTM
@@ -1321,7 +1422,7 @@ def __init__(
                 self.convlstm1 = ConvLSTM(
                     input_dim=C,
                     hidden_dim=32,
-                    kernel_size=3,
+                    kernel_size=(3, 3),
                     num_layers=1,
                     batch_first=True,
                     return_all_layers=True,
@@ -1329,7 +1430,7 @@ def __init__(
                 self.convlstm2 = ConvLSTM(
                     input_dim=32,
                     hidden_dim=32,
-                    kernel_size=3,
+                    kernel_size=(3, 3),
                     num_layers=1,
                     batch_first=True,
                     return_all_layers=True,
diff --git a/pytest.ini b/pytest.ini
new file mode 100644
index 0000000..8b6278f
--- /dev/null
+++ b/pytest.ini
@@ -0,0 +1,4 @@
+[pytest]
+testpaths = tests
+python_files = test_*.py
+addopts = -ra
\ No newline at end of file
diff --git a/tests/test_deeplearning_autoencoders.py b/tests/test_deeplearning_autoencoders.py
new file mode 100644
index 0000000..d4e7fa3
--- /dev/null
+++ b/tests/test_deeplearning_autoencoders.py
@@ -0,0 +1,241 @@
+"""
+Smoke tests for BlueMath_tk autoencoders.
+
+These tests are intentionally small so they can run quickly in CI and in a local
+Anaconda environment.
+
+Run from the repository root with:
+
+    pytest -q tests/test_deeplearning_autoencoders.py
+
+Some tests are marked xfail because they document likely current bugs in the
+implementation. After fixing each issue, remove the corresponding xfail marker.
+"""
+
+import numpy as np
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from bluemath_tk.deeplearning.autoencoders import (
+    CNNAutoencoder,
+    ConvLSTMAutoencoder,
+    HybridConvLSTMTransformerAutoencoder,
+    LSTMAutoencoder,
+    OrthogonalAutoencoder,
+    StandardAutoencoder,
+    VisionTransformerAutoencoder,
+)
+
+
+@pytest.fixture(autouse=True)
+def _set_reproducible_seed():
+    """Keep tests deterministic and avoid excessive CPU thread use."""
+    np.random.seed(123)
+    torch.manual_seed(123)
+    torch.set_num_threads(1)
+
+
+def _fit_kwargs():
+    """Common tiny training configuration."""
+    return dict(
+        epochs=1,
+        batch_size=4,
+        validation_split=0.25,
+        patience=2,
+        verbose=0,
+        learning_rate=1e-3,
+    )
+
+
+def test_standard_autoencoder_fit_predict_encode_shapes():
+    """StandardAutoencoder should reconstruct 2D tabular input."""
+    X = np.random.randn(16, 10).astype("float32")
+
+    ae = StandardAutoencoder(
+        k=3,
+        hidden_dims=[8],
+        device="cpu",
+    )
+
+    history = ae.fit(X, **_fit_kwargs())
+    X_hat = ae.predict(X, batch_size=4, verbose=0)
+    Z = ae.encode(X, batch_size=4, verbose=0)
+
+    assert set(history) == {"train_loss", "val_loss"}
+    assert len(history["train_loss"]) >= 1
+    assert len(history["val_loss"]) >= 1
+    assert X_hat.shape == X.shape
+    assert Z.shape == (16, 3)
+    assert np.isfinite(X_hat).all()
+    assert np.isfinite(Z).all()
+
+
+def test_orthogonal_autoencoder_fit_predict_encode_shapes():
+    """OrthogonalAutoencoder should reconstruct 2D tabular input."""
+    X = np.random.randn(16, 10).astype("float32")
+
+    ae = OrthogonalAutoencoder(
+        k=3,
+        hidden_dims=[8],
+        lambda_W=1e-4,
+        lambda_Z=1e-4,
+        device="cpu",
+    )
+
+    history = ae.fit(X, **_fit_kwargs())
+    X_hat = ae.predict(X, batch_size=4, verbose=0)
+    Z = ae.encode(X, batch_size=4, verbose=0)
+
+    assert set(history) == {"train_loss", "val_loss"}
+    assert X_hat.shape == X.shape
+    assert Z.shape == (16, 3)
+    assert np.isfinite(X_hat).all()
+    assert np.isfinite(Z).all()
+
+
+def test_lstm_autoencoder_fit_predict_encode_shapes():
+    """LSTMAutoencoder should reconstruct sequence input."""
+    X = np.random.randn(16, 5, 3).astype("float32")
+
+    ae = LSTMAutoencoder(
+        k=4,
+        hidden=(8, 6),
+        device="cpu",
+    )
+
+    history = ae.fit(X, **_fit_kwargs())
+    X_hat = ae.predict(X, batch_size=4, verbose=0)
+    Z = ae.encode(X, batch_size=4, verbose=0)
+
+    assert set(history) == {"train_loss", "val_loss"}
+    assert X_hat.shape == X.shape
+    assert Z.shape == (16, 4)
+    assert np.isfinite(X_hat).all()
+    assert np.isfinite(Z).all()
+
+
+def test_cnn_autoencoder_fit_predict_encode_shapes():
+    """CNNAutoencoder should reconstruct channels-first image/grid input."""
+    X = np.random.randn(16, 1, 8, 8).astype("float32")
+
+    ae = CNNAutoencoder(
+        k=4,
+        device="cpu",
+    )
+
+    history = ae.fit(X, **_fit_kwargs())
+    X_hat = ae.predict(X, batch_size=4, verbose=0)
+    Z = ae.encode(X, batch_size=4, verbose=0)
+
+    assert set(history) == {"train_loss", "val_loss"}
+    assert X_hat.shape == X.shape
+    assert Z.shape == (16, 4)
+    assert np.isfinite(X_hat).all()
+    assert np.isfinite(Z).all()
+
+
+def test_vit_autoencoder_d_model_can_differ_from_patch_dimension():
+    """
+    VisionTransformerAutoencoder should allow d_model != patch_size*patch_size*C.
+
+    For C=1 and patch_size=4, patch dimension is 16.
+    This test uses d_model=8 to catch missing decoder projection bugs.
+    """
+    X = np.random.randn(16, 1, 8, 8).astype("float32")
+
+    ae = VisionTransformerAutoencoder(
+        k=4,
+        patch_size=4,
+        d_model=8,
+        depth_enc=1,
+        depth_dec=1,
+        heads=2,
+        device="cpu",
+    )
+
+    history = ae.fit(X, **_fit_kwargs())
+    X_hat = ae.predict(X, batch_size=4, verbose=0)
+    Z = ae.encode(X, batch_size=4, verbose=0)
+
+    assert set(history) == {"train_loss", "val_loss"}
+    assert X_hat.shape == X.shape
+    assert Z.shape == (16, 4)
+    assert np.isfinite(X_hat).all()
+    assert np.isfinite(Z).all()
+
+
+def test_convlstm_autoencoder_default_fit_reconstructs_documented_single_frame():
+    """
+    ConvLSTMAutoencoder example/docstring says fit(X) should work and predict
+    returns a single reconstructed frame with shape (B, C, H, W).
+    """
+    X = np.random.randn(16, 3, 1, 8, 8).astype("float32")
+
+    ae = ConvLSTMAutoencoder(
+        k=4,
+        device="cpu",
+    )
+
+    history = ae.fit(X, **_fit_kwargs())
+    X_hat = ae.predict(X, batch_size=4, verbose=0)
+    Z = ae.encode(X, batch_size=4, verbose=0)
+
+    assert set(history) == {"train_loss", "val_loss"}
+    assert X_hat.shape == X[:, -1].shape
+    assert Z.shape == (16, 4)
+    assert np.isfinite(X_hat).all()
+    assert np.isfinite(Z).all()
+
+
+def test_hybrid_autoencoder_default_fit_reconstructs_documented_single_frame():
+    """
+    HybridConvLSTMTransformerAutoencoder example/docstring says fit(X) should
+    work and predict returns a single reconstructed frame with shape (B, C, H, W).
+    """
+    X = np.random.randn(16, 3, 1, 8, 8).astype("float32")
+
+    ae = HybridConvLSTMTransformerAutoencoder(
+        k=4,
+        d_model=8,
+        n_heads=2,
+        n_layers=1,
+        efficient_attention="linear",
+        device="cpu",
+    )
+
+    history = ae.fit(X, **_fit_kwargs())
+    X_hat = ae.predict(X, batch_size=4, verbose=0)
+    Z = ae.encode(X, batch_size=4, verbose=0)
+
+    assert set(history) == {"train_loss", "val_loss"}
+    assert X_hat.shape == X[:, -1].shape
+    assert Z.shape == (16, 4)
+    assert np.isfinite(X_hat).all()
+    assert np.isfinite(Z).all()
+
+
+def test_standard_autoencoder_multidimensional_flatten_contract():
+    """
+    Document the current ambiguity in StandardAutoencoder.
+
+    The docstring says multidimensional inputs are automatically flattened.
+    If that is intended, fit/predict should work for (B, C, H) input.
+    """
+    X = np.random.randn(16, 2, 5).astype("float32")
+
+    ae = StandardAutoencoder(
+        k=3,
+        hidden_dims=[8],
+        device="cpu",
+    )
+
+    history = ae.fit(X, **_fit_kwargs())
+    X_hat = ae.predict(X, batch_size=4, verbose=0)
+    Z = ae.encode(X, batch_size=4, verbose=0)
+
+    assert set(history) == {"train_loss", "val_loss"}
+    assert X_hat.shape == X.shape
+    assert Z.shape == (16, 3)
+    assert np.isfinite(X_hat).all()
+    assert np.isfinite(Z).all()
\ No newline at end of file