diff --git a/stable_pretraining/backbone/vit.py b/stable_pretraining/backbone/vit.py
index 921f3cd57..076428d34 100644
--- a/stable_pretraining/backbone/vit.py
+++ b/stable_pretraining/backbone/vit.py
@@ -232,9 +232,9 @@ class MaskedEncoder(nn.Module):
     Wraps a timm ViT model and adds flexible masking via :class:`PatchMasking`.
     Handles all ViT internals: patch embedding, positional embeddings, prefix
     tokens (CLS, registers), and transformer blocks.
-    :param model_or_model_name: timm model name string or pre-instantiated nn.Module
+    :param encoder_name: timm model name string or pre-instantiated nn.Module
     :param masking: PatchMasking instance. If None, no masking is applied.
-    :param pretrained: Load pretrained weights (only when model_or_model_name is str)
+    :param pretrained: Load pretrained weights (only when encoder_name is str)
     :param img_size: Override default image size
     :param patch_size: Override default patch size (will reinitialize patch_embed)
     :param dynamic_img_size: Enable dynamic image size support with pos_embed interpolation
@@ -243,7 +243,7 @@ class MaskedEncoder(nn.Module):
 
         masking = PatchMasking(mask_ratio=0.75, block_size=4)
         encoder = MaskedEncoder(
-            model_or_model_name="vit_base_patch16_224",
+            encoder_name="vit_base_patch16_224",
             masking=masking,
             pretrained=True,
         )
@@ -256,7 +256,7 @@ class MaskedEncoder(nn.Module):
 
     def __init__(
         self,
-        model_or_model_name: Union[str, nn.Module] = "vit_base_patch16_224",
+        encoder_name: Union[str, nn.Module] = "vit_base_patch16_224",
         masking: Optional[PatchMasking] = None,
         pretrained: bool = False,
         img_size: Optional[Union[int, Tuple[int, int]]] = None,
@@ -268,7 +268,7 @@ def __init__(
         self.dynamic_img_size = dynamic_img_size
         self.masking = masking
         # === Load or use provided encoder ===
-        if isinstance(model_or_model_name, str):
+        if isinstance(encoder_name, str):
             create_kwargs = {
                 "pretrained": pretrained,
                 "num_classes": 0,
@@ -286,7 +286,7 @@ def __init__(
             if norm_layer is not None:
                 create_kwargs["norm_layer"] = norm_layer
 
-            self.vit = timm.create_model(model_or_model_name, **create_kwargs)
+            self.vit = timm.create_model(encoder_name, **create_kwargs)
         else:
             logger.warning(
                 "MaskedEncoder received a pre-instantiated nn.Module. "
@@ -294,7 +294,7 @@ def __init__(
                 "patch_embed, pos_embed, cls_token, blocks, norm, etc. "
                 "If you pass a non-timm module, unexpected errors may occur."
             )
-            self.vit = model_or_model_name
+            self.vit = encoder_name
             if patch_size is not None:
                 self._rebuild_patch_embed(patch_size, img_size)
             # Remove classification head if present
diff --git a/stable_pretraining/methods/barlow_twins.py b/stable_pretraining/methods/barlow_twins.py
index 511a0ef49..6cb82b3a5 100644
--- a/stable_pretraining/methods/barlow_twins.py
+++ b/stable_pretraining/methods/barlow_twins.py
@@ -57,6 +57,9 @@ class BarlowTwins(Module):
         (default ``5.1e-3`` from the paper).
     :param low_resolution: Adapt first conv for low-res input.
     :param pretrained: Load pretrained timm weights.
+    :param embed_dim: Backbone output dimension. Inferred automatically for timm
+        models; must be provided explicitly for custom encoders that do not expose
+        an ``.embed_dim`` attribute.
     """
 
     def __init__(
@@ -66,6 +69,7 @@ def __init__(
         lambd: float = 5.1e-3,
         low_resolution: bool = False,
         pretrained: bool = False,
+        embed_dim: Optional[int] = None,
     ):
         super().__init__()
 
@@ -76,11 +80,18 @@ def __init__(
                 low_resolution=low_resolution,
                 pretrained=pretrained,
             )
+            embed_dim = self.backbone.embed_dim
         else:
             self.backbone = encoder_name
-
-        with torch.no_grad():
-            embed_dim = self.backbone(torch.zeros(1, 3, 224, 224)).shape[-1]
+            if embed_dim is None:
+                if hasattr(encoder_name, "embed_dim"):
+                    embed_dim = encoder_name.embed_dim
+                else:
+                    raise ValueError(
+                        "embed_dim must be provided when the encoder does not expose "
+                        "an .embed_dim attribute. timm models expose this automatically; "
+                        "for custom encoders, pass embed_dim explicitly."
+                    )
         self.embed_dim = embed_dim
 
         self.projector = _build_barlow_projector(embed_dim, list(projector_dims))
diff --git a/stable_pretraining/methods/byol.py b/stable_pretraining/methods/byol.py
index 0a26a3591..7048e5307 100644
--- a/stable_pretraining/methods/byol.py
+++ b/stable_pretraining/methods/byol.py
@@ -61,6 +61,9 @@ class BYOL(Module):
     :param ema_decay_end: Final EMA decay (default 1.0).
     :param low_resolution: Adapt first conv for low-res input.
     :param pretrained: Load pretrained timm weights.
+    :param embed_dim: Backbone output dimension. Inferred automatically for timm
+        models; must be provided explicitly for custom encoders that do not expose
+        an ``.embed_dim`` attribute.
 
     Note:
         Use :class:`~stable_pretraining.callbacks.TeacherStudentCallback`
@@ -76,6 +79,7 @@ def __init__(
         ema_decay_end: float = 1.0,
         low_resolution: bool = False,
         pretrained: bool = False,
+        embed_dim: Optional[int] = None,
     ):
         super().__init__()
 
@@ -86,11 +90,18 @@ def __init__(
                 low_resolution=low_resolution,
                 pretrained=pretrained,
             )
+            embed_dim = base_backbone.embed_dim
         else:
             base_backbone = encoder_name
-
-        with torch.no_grad():
-            embed_dim = base_backbone(torch.zeros(1, 3, 224, 224)).shape[-1]
+            if embed_dim is None:
+                if hasattr(encoder_name, "embed_dim"):
+                    embed_dim = encoder_name.embed_dim
+                else:
+                    raise ValueError(
+                        "embed_dim must be provided when the encoder does not expose "
+                        "an .embed_dim attribute. timm models expose this automatically; "
+                        "for custom encoders, pass embed_dim explicitly."
+                    )
         self.embed_dim = embed_dim
 
         if len(projector_dims) != 2 or len(predictor_dims) != 2:
diff --git a/stable_pretraining/methods/dino.py b/stable_pretraining/methods/dino.py
index 5c5576474..a2316d7e7 100644
--- a/stable_pretraining/methods/dino.py
+++ b/stable_pretraining/methods/dino.py
@@ -103,6 +103,9 @@ class DINO(Module):
     :param ema_decay_end: Final EMA (default 1.0).
     :param encoder_kwargs: Extra kwargs forwarded to ``timm.create_model``.
     :param pretrained: Load pretrained timm weights for the encoder.
+    :param embed_dim: Backbone output dimension. Inferred automatically for timm
+        models; must be provided explicitly for custom encoders that do not expose
+        an ``.embed_dim`` attribute.
 
     Example::
 
@@ -130,6 +133,7 @@ def __init__(
         ema_decay_end: float = 1.0,
         encoder_kwargs: Optional[dict] = None,
         pretrained: bool = False,
+        embed_dim: Optional[int] = None,
     ):
         super().__init__()
 
@@ -139,11 +143,18 @@ def __init__(
             kw = dict(num_classes=0, pretrained=pretrained)
             kw.update(encoder_kwargs or {})
             base_backbone = timm.create_model(encoder_name, **kw)
+            embed_dim = base_backbone.embed_dim
         else:
             base_backbone = encoder_name
-
-        with torch.no_grad():
-            embed_dim = _to_cls(base_backbone(torch.zeros(1, 3, 224, 224))).shape[-1]
+            if embed_dim is None:
+                if hasattr(encoder_name, "embed_dim"):
+                    embed_dim = encoder_name.embed_dim
+                else:
+                    raise ValueError(
+                        "embed_dim must be provided when the encoder does not expose "
+                        "an .embed_dim attribute. timm models expose this automatically; "
+                        "for custom encoders, pass embed_dim explicitly."
+                    )
         self.embed_dim = embed_dim
         self.n_prototypes = n_prototypes
 
diff --git a/stable_pretraining/methods/ijepa.py b/stable_pretraining/methods/ijepa.py
index 56f67dcfb..3344fbd91 100644
--- a/stable_pretraining/methods/ijepa.py
+++ b/stable_pretraining/methods/ijepa.py
@@ -16,7 +16,7 @@
 
     # Create model
     model = IJEPA(
-        model_or_model_name="vit_base_patch16_224",
+        encoder_name="vit_base_patch16_224",
         predictor_embed_dim=384,
         predictor_depth=6,
         num_targets=4,
@@ -82,7 +82,7 @@ class IJEPA(Module):
     The context encoder is wrapped with :class:`TeacherStudentWrapper`, enabling
     automatic EMA updates via :class:`TeacherStudentCallback`.
 
-    :param model_or_model_name: timm model name string or pre-instantiated nn.Module
+    :param encoder_name: timm model name string or pre-instantiated nn.Module
     :param predictor_embed_dim: Predictor hidden dimension (default: 384)
     :param predictor_depth: Number of predictor blocks (default: 6)
     :param num_targets: Number of target blocks to sample (default: 4)
@@ -141,7 +141,7 @@ def configure_optimizers(self):
 
     def __init__(
         self,
-        model_or_model_name: Union[str, nn.Module] = "vit_base_patch16_224",
+        encoder_name: Union[str, nn.Module] = "vit_base_patch16_224",
         predictor_embed_dim: int = 384,
         predictor_depth: int = 6,
         num_targets: int = 4,
@@ -156,7 +156,7 @@ def __init__(
 
         # Encoder with EMA wrapper (enables TeacherStudentCallback)
         base_encoder = MaskedEncoder(
-            model_or_model_name,
+            encoder_name,
             masking=None,
             pretrained=pretrained,
         )
diff --git a/stable_pretraining/methods/lejepa.py b/stable_pretraining/methods/lejepa.py
index 64aa54ce1..5e1f539c6 100644
--- a/stable_pretraining/methods/lejepa.py
+++ b/stable_pretraining/methods/lejepa.py
@@ -160,7 +160,7 @@ class LeJEPA(Module):
     The SIGReg term is a sliced goodness-of-fit test that pushes
     projected embeddings toward an isotropic Gaussian, averaged over views.
 
-    :param encoder_name: timm model name (e.g., ``"vit_base_patch16_224"``)
+    :param encoder_name: timm model name string or pre-instantiated nn.Module
     :param projector: Optional projection head.  When ``None``, a 3-layer
         BN+ReLU MLP (``embed_dim → 2048 → 2048 → 512``) is created.
     :param n_slices: Random projection directions for the goodness-of-fit test (default: 1024)
@@ -168,6 +168,9 @@ class LeJEPA(Module):
     :param n_points: EP quadrature nodes (default: 17)
     :param lamb: SIGReg weight λ (default: 0.02)
     :param pretrained: Load pretrained timm weights
+    :param embed_dim: Output dimension of the custom encoder. Required when
+    ``encoder_name`` is an ``nn.Module``; ignored when a timm model name
+    string is passed (inferred automatically via ``backbone.embed_dim``).
 
     Example::
 
@@ -216,18 +219,21 @@ def __init__(
         lamb: float = 0.02,
         pretrained: bool = False,
         drop_path_rate: float = 0.1,
+        embed_dim: Optional[int] = None,
     ):
         super().__init__()
-
-        self.backbone = timm.create_model(
-            encoder_name,
-            pretrained=pretrained,
-            num_classes=0,
-            **({"dynamic_img_size": True} if "vit" in encoder_name else {}),
-            drop_path_rate=drop_path_rate,
-        )
-
-        embed_dim = self.backbone.embed_dim
+        if isinstance(encoder_name, str):
+            self.backbone = timm.create_model(
+                encoder_name,
+                pretrained=pretrained,
+                num_classes=0,
+                **({"dynamic_img_size": True} if "vit" in encoder_name else {}),
+                drop_path_rate=drop_path_rate,
+            )
+            embed_dim = self.backbone.embed_dim
+        else:
+            self.backbone = encoder_name
+            embed_dim = embed_dim
 
         if projector is None:
             projector = nn.Sequential(
diff --git a/stable_pretraining/methods/mae.py b/stable_pretraining/methods/mae.py
index e7aab014c..e9c2c3f38 100644
--- a/stable_pretraining/methods/mae.py
+++ b/stable_pretraining/methods/mae.py
@@ -63,7 +63,7 @@ class MAE(Module):
         - **Decoder**: Lightweight transformer reconstructing masked patches
         - **Target**: Normalized pixel values of masked patches
 
-    :param model_or_model_name: timm model name string or pre-instantiated nn.Module
+    :param encoder_name: timm model name string or pre-instantiated nn.Module
     :param decoder_embed_dim: Decoder hidden dimension (default: 512)
     :param decoder_depth: Number of decoder blocks (default: 8)
     :param decoder_num_heads: Decoder attention heads (default: 16)
@@ -105,7 +105,7 @@ def configure_optimizers(self):
 
     def __init__(
         self,
-        model_or_model_name: Union[str, nn.Module] = "vit_base_patch16_224",
+        encoder_name: Union[str, nn.Module] = "vit_base_patch16_224",
         decoder_embed_dim: int = 512,
         decoder_depth: int = 8,
         decoder_num_heads: int = 16,
@@ -124,7 +124,7 @@ def __init__(
         else:
             self.masking = PatchMasking(mask_ratio=mask_ratio, block_size=block_size)
         self.encoder = MaskedEncoder(
-            model_or_model_name, masking=self.masking, pretrained=pretrained
+            encoder_name, masking=self.masking, pretrained=pretrained
         )
 
         embed_dim = self.encoder.embed_dim
diff --git a/stable_pretraining/methods/mocov2.py b/stable_pretraining/methods/mocov2.py
index 6e948cb68..016e7fa7e 100644
--- a/stable_pretraining/methods/mocov2.py
+++ b/stable_pretraining/methods/mocov2.py
@@ -52,6 +52,9 @@ class MoCov2(Module):
     :param ema_decay_end: Final momentum (default 0.999).
     :param low_resolution: Adapt first conv for low-res input.
     :param pretrained: Load pretrained timm weights.
+    :param embed_dim: Backbone output dimension. Inferred automatically for timm
+        models; must be provided explicitly for custom encoders that do not expose
+        an ``.embed_dim`` attribute.
     """
 
     def __init__(
@@ -64,6 +67,7 @@ def __init__(
         ema_decay_end: float = 0.999,
         low_resolution: bool = False,
         pretrained: bool = False,
+        embed_dim: Optional[int] = None,
     ):
         super().__init__()
         if isinstance(encoder_name, str):
@@ -73,11 +77,18 @@ def __init__(
                 low_resolution=low_resolution,
                 pretrained=pretrained,
             )
+            embed_dim = base.embed_dim
         else:
             base = encoder_name
-
-        with torch.no_grad():
-            embed_dim = base(torch.zeros(1, 3, 224, 224)).shape[-1]
+            if embed_dim is None:
+                if hasattr(encoder_name, "embed_dim"):
+                    embed_dim = encoder_name.embed_dim
+                else:
+                    raise ValueError(
+                        "embed_dim must be provided when the encoder does not expose "
+                        "an .embed_dim attribute. timm models expose this automatically; "
+                        "for custom encoders, pass embed_dim explicitly."
+                    )
         self.embed_dim = embed_dim
         self.temperature = temperature
 
diff --git a/stable_pretraining/methods/mocov3.py b/stable_pretraining/methods/mocov3.py
index f3c7dbb6e..b0f40abe6 100644
--- a/stable_pretraining/methods/mocov3.py
+++ b/stable_pretraining/methods/mocov3.py
@@ -79,6 +79,9 @@ class MoCov3(Module):
     :param ema_decay_end: Final EMA (default 1.0).
     :param low_resolution: Adapt first conv for low-res input.
     :param pretrained: Load pretrained timm weights.
+    :param embed_dim: Backbone output dimension. Inferred automatically for timm
+        models; must be provided explicitly for custom encoders that do not expose
+        an ``.embed_dim`` attribute.
     """
 
     def __init__(
@@ -91,6 +94,7 @@ def __init__(
         ema_decay_end: float = 1.0,
         low_resolution: bool = False,
         pretrained: bool = False,
+        embed_dim: Optional[int] = None,
     ):
         super().__init__()
 
@@ -101,11 +105,18 @@ def __init__(
                 low_resolution=low_resolution,
                 pretrained=pretrained,
             )
+            embed_dim = base.embed_dim
         else:
             base = encoder_name
-
-        with torch.no_grad():
-            embed_dim = base(torch.zeros(1, 3, 224, 224)).shape[-1]
+            if embed_dim is None:
+                if hasattr(encoder_name, "embed_dim"):
+                    embed_dim = encoder_name.embed_dim
+                else:
+                    raise ValueError(
+                        "embed_dim must be provided when the encoder does not expose "
+                        "an .embed_dim attribute. timm models expose this automatically; "
+                        "for custom encoders, pass embed_dim explicitly."
+                    )
         self.embed_dim = embed_dim
         self.temperature = temperature
 
diff --git a/stable_pretraining/methods/nnclr.py b/stable_pretraining/methods/nnclr.py
index f0405afe9..cb717c224 100644
--- a/stable_pretraining/methods/nnclr.py
+++ b/stable_pretraining/methods/nnclr.py
@@ -77,6 +77,9 @@ class NNCLR(Module):
     :param temperature: NT-Xent temperature (default 0.1).
     :param low_resolution: Adapt first conv for low-res input.
     :param pretrained: Load pretrained timm weights.
+    :param embed_dim: Backbone output dimension. Inferred automatically for timm
+        models; must be provided explicitly for custom encoders that do not expose
+        an ``.embed_dim`` attribute.
     """
 
     def __init__(
@@ -88,6 +91,7 @@ def __init__(
         temperature: float = 0.1,
         low_resolution: bool = False,
         pretrained: bool = False,
+        embed_dim: Optional[int] = None,
     ):
         super().__init__()
         if isinstance(encoder_name, str):
@@ -97,11 +101,18 @@ def __init__(
                 low_resolution=low_resolution,
                 pretrained=pretrained,
             )
+            embed_dim = self.backbone.embed_dim
         else:
             self.backbone = encoder_name
-
-        with torch.no_grad():
-            embed_dim = self.backbone(torch.zeros(1, 3, 224, 224)).shape[-1]
+            if embed_dim is None:
+                if hasattr(encoder_name, "embed_dim"):
+                    embed_dim = encoder_name.embed_dim
+                else:
+                    raise ValueError(
+                        "embed_dim must be provided when the encoder does not expose "
+                        "an .embed_dim attribute. timm models expose this automatically; "
+                        "for custom encoders, pass embed_dim explicitly."
+                    )
         self.embed_dim = embed_dim
 
         proj_hidden, proj_out = projector_dims
diff --git a/stable_pretraining/methods/pirl.py b/stable_pretraining/methods/pirl.py
index 336d20ad1..34c053785 100644
--- a/stable_pretraining/methods/pirl.py
+++ b/stable_pretraining/methods/pirl.py
@@ -75,6 +75,9 @@ class PIRL(Module):
     :param jigsaw_grid: Grid size for the jigsaw transform (default 3).
     :param low_resolution: Adapt first conv for low-res input.
     :param pretrained: Load pretrained timm weights.
+    :param embed_dim: Backbone output dimension. Inferred automatically for timm
+        models; must be provided explicitly for custom encoders that do not expose
+        an ``.embed_dim`` attribute.
     """
 
     def __init__(
@@ -87,6 +90,7 @@ def __init__(
         jigsaw_grid: int = 4,
         low_resolution: bool = False,
         pretrained: bool = False,
+        embed_dim: Optional[int] = None,
     ):
         super().__init__()
         if isinstance(encoder_name, str):
@@ -96,11 +100,18 @@ def __init__(
                 low_resolution=low_resolution,
                 pretrained=pretrained,
             )
+            embed_dim = self.backbone.embed_dim
         else:
             self.backbone = encoder_name
-
-        with torch.no_grad():
-            embed_dim = self.backbone(torch.zeros(1, 3, 224, 224)).shape[-1]
+            if embed_dim is None:
+                if hasattr(encoder_name, "embed_dim"):
+                    embed_dim = encoder_name.embed_dim
+                else:
+                    raise ValueError(
+                        "embed_dim must be provided when the encoder does not expose "
+                        "an .embed_dim attribute. timm models expose this automatically; "
+                        "for custom encoders, pass embed_dim explicitly."
+                    )
         self.embed_dim = embed_dim
         self.temperature = temperature
         self.lambda_pirl = lambda_pirl
diff --git a/stable_pretraining/methods/simclr.py b/stable_pretraining/methods/simclr.py
index bce7e59bd..232487b22 100644
--- a/stable_pretraining/methods/simclr.py
+++ b/stable_pretraining/methods/simclr.py
@@ -92,6 +92,9 @@ class SimCLR(Module):
         is common for harder/larger batches).
     :param low_resolution: Adapt first conv for 32x32 inputs (CIFAR-style).
     :param pretrained: Load pretrained timm weights for the encoder.
+    :param embed_dim: Backbone output dimension. Inferred automatically for timm
+        models; must be provided explicitly for custom encoders that do not expose
+        an ``.embed_dim`` attribute.
 
     Example::
 
@@ -119,6 +122,7 @@ def __init__(
         temperature: float = 0.5,
         low_resolution: bool = False,
         pretrained: bool = False,
+        embed_dim: Optional[int] = None,
     ):
         super().__init__()
 
@@ -129,13 +133,18 @@ def __init__(
                 low_resolution=low_resolution,
                 pretrained=pretrained,
             )
+            embed_dim = self.backbone.embed_dim
         else:
             self.backbone = encoder_name
-
-        # Detect embedding dimension by running a tiny dummy input
-        with torch.no_grad():
-            dummy = torch.zeros(1, 3, 224, 224)
-            embed_dim = self.backbone(dummy).shape[-1]
+            if embed_dim is None:
+                if hasattr(encoder_name, "embed_dim"):
+                    embed_dim = encoder_name.embed_dim
+                else:
+                    raise ValueError(
+                        "embed_dim must be provided when the encoder does not expose "
+                        "an .embed_dim attribute. timm models expose this automatically; "
+                        "for custom encoders, pass embed_dim explicitly."
+                    )
         self.embed_dim = embed_dim
 
         self.projector = _build_projector(embed_dim, list(projector_dims))
diff --git a/stable_pretraining/methods/simsiam.py b/stable_pretraining/methods/simsiam.py
index bb95314e5..c73f8e707 100644
--- a/stable_pretraining/methods/simsiam.py
+++ b/stable_pretraining/methods/simsiam.py
@@ -73,6 +73,9 @@ class SimSiam(Module):
     :param predictor_hidden_dim: Predictor bottleneck dim (default 512).
     :param low_resolution: Adapt first conv for 32x32.
     :param pretrained: Load pretrained timm weights.
+    :param embed_dim: Backbone output dimension. Inferred automatically for timm
+        models; must be provided explicitly for custom encoders that do not expose
+        an ``.embed_dim`` attribute.
     """
 
     def __init__(
@@ -82,6 +85,7 @@ def __init__(
         predictor_hidden_dim: int = 512,
         low_resolution: bool = False,
         pretrained: bool = False,
+        embed_dim: Optional[int] = None,
     ):
         super().__init__()
         if isinstance(encoder_name, str):
@@ -91,11 +95,18 @@ def __init__(
                 low_resolution=low_resolution,
                 pretrained=pretrained,
             )
+            embed_dim = self.backbone.embed_dim
         else:
             self.backbone = encoder_name
-
-        with torch.no_grad():
-            embed_dim = self.backbone(torch.zeros(1, 3, 224, 224)).shape[-1]
+            if embed_dim is None:
+                if hasattr(encoder_name, "embed_dim"):
+                    embed_dim = encoder_name.embed_dim
+                else:
+                    raise ValueError(
+                        "embed_dim must be provided when the encoder does not expose "
+                        "an .embed_dim attribute. timm models expose this automatically; "
+                        "for custom encoders, pass embed_dim explicitly."
+                    )
         self.embed_dim = embed_dim
 
         self.projector = _projector(embed_dim, projector_dim, projector_dim)
diff --git a/stable_pretraining/methods/swav.py b/stable_pretraining/methods/swav.py
index 79bfcb5ed..99e76f5af 100644
--- a/stable_pretraining/methods/swav.py
+++ b/stable_pretraining/methods/swav.py
@@ -55,6 +55,9 @@ class SwAV(Module):
     :param epsilon: Sinkhorn entropy coefficient (default 0.05).
     :param low_resolution: Adapt first conv for low-res input.
     :param pretrained: Load pretrained timm weights.
+    :param embed_dim: Backbone output dimension. Inferred automatically for timm
+        models; must be provided explicitly for custom encoders that do not expose
+        an ``.embed_dim`` attribute.
     """
 
     def __init__(
@@ -68,6 +71,7 @@ def __init__(
         low_resolution: bool = False,
         pretrained: bool = False,
         dynamic_img_size: bool = True,
+        embed_dim: Optional[int] = None,
     ):
         super().__init__()
         if isinstance(encoder_name, str):
@@ -78,11 +82,18 @@ def __init__(
                 pretrained=pretrained,
                 dynamic_img_size=dynamic_img_size,
             )
+            embed_dim = self.backbone.embed_dim
         else:
             self.backbone = encoder_name
-
-        with torch.no_grad():
-            embed_dim = self.backbone(torch.zeros(1, 3, 224, 224)).shape[-1]
+            if embed_dim is None:
+                if hasattr(encoder_name, "embed_dim"):
+                    embed_dim = encoder_name.embed_dim
+                else:
+                    raise ValueError(
+                        "embed_dim must be provided when the encoder does not expose "
+                        "an .embed_dim attribute. timm models expose this automatically; "
+                        "for custom encoders, pass embed_dim explicitly."
+                    )
         self.embed_dim = embed_dim
 
         proj_hidden, proj_out = projector_dims
diff --git a/stable_pretraining/methods/tico.py b/stable_pretraining/methods/tico.py
index f090e4b9f..7be1b047d 100644
--- a/stable_pretraining/methods/tico.py
+++ b/stable_pretraining/methods/tico.py
@@ -57,6 +57,9 @@ class TiCO(Module):
         paper used a sweep around 16-20).
     :param low_resolution: Adapt first conv for low-res input.
     :param pretrained: Load pretrained timm weights.
+    :param embed_dim: Backbone output dimension. Inferred automatically for timm
+        models; must be provided explicitly for custom encoders that do not expose
+        an ``.embed_dim`` attribute.
     """
 
     def __init__(
@@ -67,6 +70,7 @@ def __init__(
         rho: float = 20.0,
         low_resolution: bool = False,
         pretrained: bool = False,
+        embed_dim: Optional[int] = None,
     ):
         super().__init__()
         if isinstance(encoder_name, str):
@@ -76,11 +80,18 @@ def __init__(
                 low_resolution=low_resolution,
                 pretrained=pretrained,
             )
+            embed_dim = self.backbone.embed_dim
         else:
             self.backbone = encoder_name
-
-        with torch.no_grad():
-            embed_dim = self.backbone(torch.zeros(1, 3, 224, 224)).shape[-1]
+            if embed_dim is None:
+                if hasattr(encoder_name, "embed_dim"):
+                    embed_dim = encoder_name.embed_dim
+                else:
+                    raise ValueError(
+                        "embed_dim must be provided when the encoder does not expose "
+                        "an .embed_dim attribute. timm models expose this automatically; "
+                        "for custom encoders, pass embed_dim explicitly."
+                    )
         self.embed_dim = embed_dim
         self.beta = beta
         self.rho = rho
diff --git a/stable_pretraining/methods/vicreg.py b/stable_pretraining/methods/vicreg.py
index c287a728e..265c4a400 100644
--- a/stable_pretraining/methods/vicreg.py
+++ b/stable_pretraining/methods/vicreg.py
@@ -70,6 +70,9 @@ class VICReg(Module):
     :param cov_coeff: Covariance term weight (default 1.0).
     :param low_resolution: Adapt first conv for 32x32 inputs (CIFAR-style).
     :param pretrained: Load pretrained timm weights for the encoder.
+    :param embed_dim: Backbone output dimension. Inferred automatically for timm
+        models; must be provided explicitly for custom encoders that do not expose
+        an ``.embed_dim`` attribute.
     """
 
     def __init__(
@@ -81,6 +84,7 @@ def __init__(
         cov_coeff: float = 1.0,
         low_resolution: bool = False,
         pretrained: bool = False,
+        embed_dim: Optional[int] = None,
     ):
         super().__init__()
 
@@ -91,11 +95,18 @@ def __init__(
                 low_resolution=low_resolution,
                 pretrained=pretrained,
             )
+            embed_dim = self.backbone.embed_dim
         else:
             self.backbone = encoder_name
-
-        with torch.no_grad():
-            embed_dim = self.backbone(torch.zeros(1, 3, 224, 224)).shape[-1]
+            if embed_dim is None:
+                if hasattr(encoder_name, "embed_dim"):
+                    embed_dim = encoder_name.embed_dim
+                else:
+                    raise ValueError(
+                        "embed_dim must be provided when the encoder does not expose "
+                        "an .embed_dim attribute. timm models expose this automatically; "
+                        "for custom encoders, pass embed_dim explicitly."
+                    )
         self.embed_dim = embed_dim
 
         self.projector = _build_vicreg_projector(embed_dim, list(projector_dims))
diff --git a/stable_pretraining/methods/wmse.py b/stable_pretraining/methods/wmse.py
index 5c4c45f93..4852e2ddc 100644
--- a/stable_pretraining/methods/wmse.py
+++ b/stable_pretraining/methods/wmse.py
@@ -72,6 +72,9 @@ class WMSE(Module):
     :param eps: Cholesky regularisation (default ``1e-3``).
     :param low_resolution: Adapt first conv for low-res input.
     :param pretrained: Load pretrained timm weights.
+    :param embed_dim: Backbone output dimension. Inferred automatically for timm
+        models; must be provided explicitly for custom encoders that do not expose
+        an ``.embed_dim`` attribute.
     """
 
     def __init__(
@@ -81,6 +84,7 @@ def __init__(
         eps: float = 1e-3,
         low_resolution: bool = False,
         pretrained: bool = False,
+        embed_dim: Optional[int] = None,
     ):
         super().__init__()
         if isinstance(encoder_name, str):
@@ -90,11 +94,18 @@ def __init__(
                 low_resolution=low_resolution,
                 pretrained=pretrained,
             )
+            embed_dim = self.backbone.embed_dim
         else:
             self.backbone = encoder_name
-
-        with torch.no_grad():
-            embed_dim = self.backbone(torch.zeros(1, 3, 224, 224)).shape[-1]
+            if embed_dim is None:
+                if hasattr(encoder_name, "embed_dim"):
+                    embed_dim = encoder_name.embed_dim
+                else:
+                    raise ValueError(
+                        "embed_dim must be provided when the encoder does not expose "
+                        "an .embed_dim attribute. timm models expose this automatically; "
+                        "for custom encoders, pass embed_dim explicitly."
+                    )
         self.embed_dim = embed_dim
 
         proj_hidden, proj_out = projector_dims
diff --git a/stable_pretraining/tests/integration/test_ijepa_inet10.py b/stable_pretraining/tests/integration/test_ijepa_inet10.py
index aebcf9357..85ce6fea8 100644
--- a/stable_pretraining/tests/integration/test_ijepa_inet10.py
+++ b/stable_pretraining/tests/integration/test_ijepa_inet10.py
@@ -82,7 +82,7 @@ def ijepa_forward(self, batch, stage):
 
         # Create IJEPA module with vit_tiny for fast CPU testing
         module = IJEPA(
-            model_or_model_name="vit_tiny_patch16_224",
+            encoder_name="vit_tiny_patch16_224",
             predictor_embed_dim=192,
             predictor_depth=6,
             num_targets=4,
diff --git a/stable_pretraining/tests/integration/test_ijepa_module_input.py b/stable_pretraining/tests/integration/test_ijepa_module_input.py
index 9a2fd2f19..7bf59c6b0 100644
--- a/stable_pretraining/tests/integration/test_ijepa_module_input.py
+++ b/stable_pretraining/tests/integration/test_ijepa_module_input.py
@@ -31,7 +31,7 @@ class TestIJEPAModuleInput:
     """Run IJEPA with a pre-loaded backbone on imagenette for 3 steps.
 
     Mirrors test_ijepa_inet10.py but passes an nn.Module instead of a string
-    to IJEPA(), verifying the model_or_model_name interface.
+    to IJEPA(), verifying the encoder_name interface.
     """
 
     def test_ijepa_3_steps_with_loaded_backbone(self):
@@ -99,7 +99,7 @@ def ijepa_forward(self, batch, stage):
         backbone = load_backbone("vit_tiny_patch16_224", pretrained=False)
 
         module = IJEPA(
-            model_or_model_name=backbone,
+            encoder_name=backbone,
             predictor_embed_dim=192,
             predictor_depth=6,
             num_targets=4,
diff --git a/stable_pretraining/tests/integration/test_mae_inet10.py b/stable_pretraining/tests/integration/test_mae_inet10.py
index f2ad8be79..96e89fdc7 100644
--- a/stable_pretraining/tests/integration/test_mae_inet10.py
+++ b/stable_pretraining/tests/integration/test_mae_inet10.py
@@ -81,7 +81,7 @@ def mae_forward(self, batch, stage):
 
         # Create MAE module with vit_tiny for fast CPU testing
         module = MAE(
-            model_or_model_name="vit_tiny_patch16_224",
+            encoder_name="vit_tiny_patch16_224",
             decoder_embed_dim=192,
             decoder_depth=4,
             decoder_num_heads=3,
diff --git a/stable_pretraining/tests/integration/test_mae_module_input.py b/stable_pretraining/tests/integration/test_mae_module_input.py
index 5479aa5cb..65e61d428 100644
--- a/stable_pretraining/tests/integration/test_mae_module_input.py
+++ b/stable_pretraining/tests/integration/test_mae_module_input.py
@@ -31,7 +31,7 @@ class TestMAEImagenet10LoadModule:
     """Run MAE with a pre-loaded backbone on imagenette for 3 steps.
 
     Mirrors test_mae_inet10.py but passes an nn.Module instead of a string
-    to MAE(), verifying the model_or_model_name interface.
+    to MAE(), verifying the encoder_name interface.
     """
 
     def test_mae_10_steps_with_loaded_backbone(self):
@@ -98,7 +98,7 @@ def mae_forward(self, batch, stage):
         backbone = load_backbone("vit_tiny_patch16_224", pretrained=False)
 
         module = MAE(
-            model_or_model_name=backbone,
+            encoder_name=backbone,
             decoder_embed_dim=192,
             decoder_depth=4,
             decoder_num_heads=3,
diff --git a/stable_pretraining/tests/unit/test_hf_models.py b/stable_pretraining/tests/unit/test_hf_models.py
index 785bee690..d8603afa1 100644
--- a/stable_pretraining/tests/unit/test_hf_models.py
+++ b/stable_pretraining/tests/unit/test_hf_models.py
@@ -52,7 +52,7 @@ class MockSPTSystem(spt.Module):
     def __init__(self, cfg):
         super().__init__()
         self.encoder = spt.backbone.MaskedEncoder(
-            model_or_model_name=cfg.get("model_name", "vit_tiny_patch16_224"),
+            encoder_name=cfg.get("model_name", "vit_tiny_patch16_224"),
             masking=spt.backbone.PatchMasking(mask_ratio=0.5, block_size=1),
             img_size=(224, 224),
             pretrained=False,
diff --git a/stable_pretraining/tests/unit/test_methods_smoke.py b/stable_pretraining/tests/unit/test_methods_smoke.py
index 146441878..13936a855 100644
--- a/stable_pretraining/tests/unit/test_methods_smoke.py
+++ b/stable_pretraining/tests/unit/test_methods_smoke.py
@@ -205,7 +205,7 @@ def test_swav_forward_backward_multicrop():
 
 def test_ijepa_forward_backward():
     model = M.IJEPA(
-        model_or_model_name=TINY_VIT,
+        encoder_name=TINY_VIT,
         predictor_embed_dim=64,
         predictor_depth=2,
         num_targets=2,
@@ -218,7 +218,7 @@ def test_ijepa_forward_backward():
 
 def test_mae_forward_backward():
     model = M.MAE(
-        model_or_model_name=TINY_VIT,
+        encoder_name=TINY_VIT,
         decoder_embed_dim=64,
         decoder_depth=2,
         decoder_num_heads=2,