diff --git a/tsimcne/tsimcne.py b/tsimcne/tsimcne.py
index 7936b79..7da1411 100644
--- a/tsimcne/tsimcne.py
+++ b/tsimcne/tsimcne.py
@@ -475,7 +475,6 @@ def fit(self, X: torch.utils.data.Dataset | str):
                     model=model,
                     loss=p.loss,
                     metric=p.metric,
-                    out_dim=self.out_dim,
                     **train_kwargs,
                 )
             trainer = pl.Trainer(max_epochs=n_epochs, devices=1)