diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index db9738e6e7..0b46baea79 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -933,6 +933,8 @@ def process_weights_after_loading(self, layer): ) layer.w13_weight.is_shuffled = True layer.w2_weight.is_shuffled = True + layer.w13_weight.is_guinterleave = self.is_guinterleave + layer.w2_weight.is_guinterleave = self.is_guinterleave # shuffle scale w13_scale_2d = layer.w13_weight_scale.reshape(