diff --git a/train.py b/train.py
index c9203f1..a1f8ee1 100644
--- a/train.py
+++ b/train.py
@@ -155,10 +155,15 @@ def train(config_path):
     # Manually wrap model with FSDP before passing to trainer
     from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
     from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+    from functools import partial
     
     def get_auto_wrap_policy(model):
         from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeDecoderLayer
-        return transformer_auto_wrap_policy({Qwen3_5MoeDecoderLayer}, model)
+        return partial(
+            transformer_auto_wrap_policy,
+            transformer_layer_cls={Qwen3_5MoeDecoderLayer},
+            min_num_params=1,
+        )
     
     # Wrap model with FSDP on CPU first
     print("Wrapping model with FSDP on CPU...")