diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py
index d08e2cee..3aa43280 100644
--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -317,7 +317,7 @@ def tensor_auto_parallel(
     elif hasattr(model, "shard"):
         try:
             model.shard(group)  # type: ignore
-            return patch_distributed_model(model)
+            return model
         except (AttributeError, TypeError, NameError):
             pass