From 4dcdca7b714e8beac973b7287010189daaef3866 Mon Sep 17 00:00:00 2001
From: Ryuichi Leo Takashige <leo@exolabs.net>
Date: Tue, 20 Jan 2026 14:44:40 +0000
Subject: [PATCH] Fix?

---
 src/exo/worker/engines/mlx/auto_parallel.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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