diff --git a/network/body/embedding.py b/network/body/embedding.py
index 83db476..5e618ba 100644
--- a/network/body/embedding.py
+++ b/network/body/embedding.py
@@ -263,7 +263,11 @@ class PETBody(nn.Module):
     def __init__(
             self, num_feat, num_keep, feature_drop, projection_dim, local, K, num_local,
             num_layers, num_heads, drop_probability, talking_head, layer_scale,
-            layer_scale_init, dropout, mode, use_adapter: bool = False
+            layer_scale_init, dropout, mode, use_adapter: bool = False,
+            use_moe: bool = False, moe_base_num_experts: int = 4,
+            moe_base_select_top_k: int = 2, moe_num_shared_experts: int = 0,
+            moe_expert_segmentation_factor: int = 1, moe_scale_expert_dim: bool = False,
+            moe_alpha: float = 0.01, moe_cz: float = 0.0, moe_use_router_noise: bool = False
     ):
         super().__init__()
         self.num_keep = num_keep
@@ -295,7 +299,12 @@ def __init__(
         self.transformer_blocks = nn.ModuleList([
             TransformerBlockModule(
                 projection_dim, num_heads, dropout, talking_head, layer_scale, layer_scale_init,
-                drop_probability
+                drop_probability, use_moe=use_moe, moe_base_num_experts=moe_base_num_experts,
+                moe_base_select_top_k=moe_base_select_top_k, moe_num_shared_experts=moe_num_shared_experts,
+                moe_expert_segmentation_factor=moe_expert_segmentation_factor,
+                moe_scale_expert_dim=moe_scale_expert_dim,
+                moe_alpha=moe_alpha, moe_cz=moe_cz,
+                moe_use_router_noise=moe_use_router_noise
             )
             for _ in range(num_layers)
         ])
@@ -350,16 +359,21 @@ def forward(self,
             encoded = local_features + encoded  # Combine with original features
 
         skip_connection = encoded
+        moe_l_aux = encoded.new_zeros(())
+        moe_cz_lz = encoded.new_zeros(())
         for itransformer, transformer_block in enumerate(self.transformer_blocks):
             encoded = transformer_block(
                 x=encoded,
                 mask=mask,
                 attn_mask=attn_mask
             )
+            moe_l_aux += transformer_block.moe_l_aux.to(encoded.device)
+            moe_cz_lz += transformer_block.moe_cz_lz.to(encoded.device)
             if self.use_adapter:
                 encoded = self.adapters[itransformer](encoded)
                 encoded = encoded * mask.float()
-
+        self.moe_l_aux = moe_l_aux
+        self.moe_cz_lz = moe_cz_lz
 
         return torch.add(encoded, skip_connection)
 
diff --git a/network/evenet_model.py b/network/evenet_model.py
index 016ab3f..068c0f6 100644
--- a/network/evenet_model.py
+++ b/network/evenet_model.py
@@ -21,6 +21,7 @@
 from torch import Tensor, nn
 from typing import Dict, Optional, Any, Union
 import re
+import logging
 
 
 class EveNetModel(nn.Module):
@@ -155,6 +156,15 @@ def __init__(
             layer_scale_init=pet_config.layer_scale_init,
             dropout=pet_config.dropout,
             mode=pet_config.mode,
+            use_moe=pet_config.use_moe,
+            moe_base_num_experts=pet_config.moe_base_num_experts,
+            moe_base_select_top_k=pet_config.moe_base_select_top_k,
+            moe_num_shared_experts=pet_config.moe_num_shared_experts,
+            moe_expert_segmentation_factor=pet_config.moe_expert_segmentation_factor,
+            moe_scale_expert_dim=pet_config.moe_scale_expert_dim,
+            moe_alpha=pet_config.moe_alpha,
+            moe_cz=pet_config.moe_cz,
+            moe_use_router_noise=pet_config.moe_use_router_noise,
         )
 
         # [2] Classification + Regression + Assignment Body
@@ -317,6 +327,41 @@ def __init__(
             ("deterministic", self.include_classification or self.include_assignment or self.include_regression or self.include_segmentation),
         ]
 
+        self._log_backbone_setup()
+
+    def _log_backbone_setup(self) -> None:
+        logger = logging.getLogger(__name__)
+        pet_cfg = self.network_cfg.Body.PET
+
+        pretrain_path = getattr(getattr(self.options, "Training", None), "pretrain_model_load_path", None)
+        if pretrain_path:
+            logger.info(f"[Backbone] Pretrain path       : {pretrain_path}")
+        else:
+            logger.warning("[Backbone] No pretrain_model_load_path set — training from scratch.")
+
+        logger.info(
+            f"[Backbone] PET config          : "
+            f"layers={pet_cfg.num_layers}, "
+            f"heads={pet_cfg.num_heads}, "
+            f"dim={pet_cfg.hidden_dim}, "
+            f"mode={pet_cfg.mode}"
+        )
+        logger.info(
+            f"[Backbone] MoE enabled         : {pet_cfg.use_moe}"
+        )
+        if pet_cfg.use_moe:
+            logger.info(
+                f"[Backbone] MoE config          : "
+                f"num_experts={pet_cfg.moe_base_num_experts}, "
+                f"top_k={pet_cfg.moe_base_select_top_k}, "
+                f"shared_experts={pet_cfg.moe_num_shared_experts}, "
+                f"seg_factor={pet_cfg.moe_expert_segmentation_factor}, "
+                f"scale_dim={pet_cfg.moe_scale_expert_dim}, "
+                f"alpha={pet_cfg.moe_alpha}, "
+                f"cz={pet_cfg.moe_cz}, "
+                f"router_noise={pet_cfg.moe_use_router_noise}"
+            )
+
     def forward(
             self, x: Dict[str, Tensor], time: Tensor,
             progressive_params: dict = None,
@@ -461,6 +506,8 @@ def forward(
 
         full_input_point_cloud = None
         full_global_conditions = None
+        moe_l_aux_total = torch.zeros((), device=input_point_cloud.device, dtype=input_point_cloud.dtype)
+        moe_cz_lz_total = torch.zeros((), device=input_point_cloud.device, dtype=input_point_cloud.dtype)
 
         for schedule_name, flag  in schedules:
             if not flag:
@@ -537,6 +584,8 @@ def forward(
                 time=full_time,
                 time_masking=time_masking
             )
+            moe_l_aux_total += self.PET.moe_l_aux
+            moe_cz_lz_total += self.PET.moe_cz_lz
 
             if schedule_name == "deterministic" or schedule_name == "generation":
                 ######################################
@@ -655,7 +704,9 @@ def forward(
             # "full_global_conditions": full_global_conditions,
             "alpha": alpha,
             "segmentation-mask": outputs.get("deterministic", {}).get("segmentation-out", {}).get("pred_masks", None),
-            "segmentation-aux": outputs.get("deterministic", {}).get("segmentation-out", {}).get("aux_outputs", None)
+            "segmentation-aux": outputs.get("deterministic", {}).get("segmentation-out", {}).get("aux_outputs", None),
+            "L_aux": moe_l_aux_total,
+            "cz_Lz": moe_cz_lz_total,
         }
 
     def predict_diffusion_vector(
diff --git a/network/layers/transformer.py b/network/layers/transformer.py
index b967084..bd4372e 100644
--- a/network/layers/transformer.py
+++ b/network/layers/transformer.py
@@ -1,6 +1,9 @@
+import logging
+
 import torch.nn as nn
 from torch import Tensor
 import torch
+import torch.nn.functional as F
 
 from evenet.network.layers.utils import TalkingHeadAttention, StochasticDepth, LayerScale
 from evenet.network.layers.linear_block import GRUGate, GRUBlock
@@ -8,9 +11,234 @@
 
 from typing import Optional
 
+_moe_logger = logging.getLogger(__name__)
+
+
+class Gate(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        num_experts: int,
+        select_top_k: int,
+        use_router_noise: bool
+    ) -> None:
+        super().__init__()
+        self.router = nn.Linear(embed_dim, num_experts, bias=False)
+        self.noise_router = nn.Linear(embed_dim, num_experts, bias=False) if use_router_noise else None
+        self.select_top_k = select_top_k
+
+    def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+        router_logits = self.router(x)
+        noisy_router_logits = router_logits
+
+        # noisy top-k routing during training to keep exploration healthy
+        if self.training and self.noise_router is not None:
+            noise_std = F.softplus(self.noise_router(x))
+            noisy_router_logits = noisy_router_logits + torch.randn_like(noisy_router_logits) * noise_std
+
+        # get top-k expert scores per object
+        topk_logits, topk_indices = torch.topk(noisy_router_logits, self.select_top_k, dim=-1)
+        # probability distribution over selected experts only
+        topk_weights = torch.softmax(topk_logits, dim=-1)
+
+        # dense gate weights over all experts from clean logits; used for losses/stats
+        dense_gate_weights = torch.softmax(router_logits, dim=-1)
+
+        return router_logits, dense_gate_weights, topk_weights, topk_indices
+
+
+class Expert(nn.Module):
+    def __init__(self, embed_dim: int, feedforward_dim: int, dropout: float) -> None:
+        super().__init__()
+        self.ffn = nn.Sequential(
+            nn.Linear(embed_dim, feedforward_dim),
+            # GELU activation maintained from original PET FFN
+            nn.GELU(approximate="none"),
+            nn.Dropout(dropout),
+            nn.Linear(feedforward_dim, embed_dim),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.ffn(x)
+
+
+class MoE(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        feedforward_dim: int,
+        base_num_experts: int,
+        base_select_top_k: int,
+        num_shared_experts: int,
+        expert_segmentation_factor: int,
+        scale_expert_dim: bool,
+        alpha: float,
+        c_z: float,
+        use_router_noise: bool,
+        dropout: float
+    ) -> None:
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.base_num_experts = base_num_experts
+        self.base_select_top_k = base_select_top_k
+        self.expert_segmentation_factor = expert_segmentation_factor
+        self.num_shared_experts = num_shared_experts
+        self.alpha = alpha
+        self.c_z = c_z
+
+        total_experts = self.base_num_experts * self.expert_segmentation_factor
+        # num_experts is the total budget - routed experts fill the remainder after reserving shared slots
+        self.num_experts = total_experts - num_shared_experts
+        self.select_top_k = self.base_select_top_k * self.expert_segmentation_factor
+        # when scale_expert_dim is True divide each expert's hidden dim by select_top_k to keep per-token compute constant vs. a vanilla FFN
+        # note: even shared experts are being impacted by segmentation scaling of k
+        self.expert_hidden_dim = int(feedforward_dim / (self.select_top_k + self.num_shared_experts)) if scale_expert_dim else feedforward_dim
+
+        self.gate = Gate(embed_dim, self.num_experts, self.select_top_k, use_router_noise=use_router_noise)
+        self.routed_experts = nn.ModuleList([
+                Expert(embed_dim, self.expert_hidden_dim, dropout)
+                for _ in range(self.num_experts)
+        ])
+        self.shared_experts = nn.ModuleList([
+                Expert(embed_dim, self.expert_hidden_dim, dropout)
+                for _ in range(self.num_shared_experts)
+        ])
+        self._forward_logged = False
+        self.register_buffer('expert_dispatch_counts', torch.zeros(self.num_experts, dtype=torch.long))
+
+    def reset_expert_dispatch_counts(self) -> None:
+        """Reset the accumulated eval-time dispatch counts to zero."""
+        self.expert_dispatch_counts.zero_()
+
+    def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]:
+        original_shape = x.shape
+        # collapse batch/object axes to a 2D tensor so that each row corresponds to a single object to route to experts
+        x = x.reshape(-1, x.shape[-1])
+        num_objects = x.shape[0]
+
+        router_logits, dense_gate_weights, topk_weights, topk_indices = self.gate(x)
+
+        if not self._forward_logged:
+            expert_w_in  = self.routed_experts[0].ffn[0].weight.shape  # (hidden, embed)
+            expert_w_out = self.routed_experts[0].ffn[3].weight.shape  # (embed, hidden)
+            gate_w       = self.gate.router.weight.shape                # (num_experts, embed)
+
+            shape_ok = (
+                x.shape[-1]             == self.embed_dim
+                and router_logits.shape == (num_objects, self.num_experts)
+                and topk_indices.shape  == (num_objects, self.select_top_k)
+                and expert_w_in[0]      == self.expert_hidden_dim
+                and expert_w_in[1]      == self.embed_dim
+                and gate_w[0]           == self.num_experts
+            )
+            print(
+                f"[MoE] First forward — input: {tuple(original_shape)}  "
+                f"(flattened tokens: {num_objects})"
+            )
+            print(
+                f"[MoE]   gate    weight : {tuple(gate_w)}  "
+                f"→ router_logits: {tuple(router_logits.shape)}  "
+                f"topk_indices: {tuple(topk_indices.shape)}"
+            )
+            print(
+                f"[MoE]   expert  ffn[0] : {tuple(expert_w_in)}  "
+                f"ffn[3]: {tuple(expert_w_out)}  "
+                f"(expected hidden={self.expert_hidden_dim}, embed={self.embed_dim})"
+            )
+            if self.num_shared_experts > 0:
+                sh_w_in  = self.shared_experts[0].ffn[0].weight.shape
+                sh_w_out = self.shared_experts[0].ffn[3].weight.shape
+                print(
+                    f"[MoE]   shared  ffn[0] : {tuple(sh_w_in)}  "
+                    f"ffn[3]: {tuple(sh_w_out)}"
+                )
+            if not shape_ok:
+                print("[MoE]   ❌ Shape mismatch detected — check config vs checkpoint.")
+            else:
+                print("[MoE]   ✅ All shapes consistent.")
+            self._forward_logged = True
+
+        routed_output = torch.zeros((num_objects, self.embed_dim), dtype=x.dtype, device=x.device)
+        objects_per_expert = torch.zeros(self.num_experts, dtype=torch.long, device=x.device)
+
+        if num_objects > 0:
+            # get flat list of each object id repeated for each of its top-k experts - e.g., [0, 0, 1, 1, 2, 2, ...]
+            object_indices = torch.arange(num_objects, device=x.device).unsqueeze(1).expand(-1, self.select_top_k).reshape(-1)
+            # get flat list of which expert each object is assigned to - of size [num_objects * top_k]
+            expert_indices = topk_indices.reshape(-1)
+            # get flat list of corresponding expert weights for each object - of size [num_objects * top_k]
+            expert_weights = topk_weights.reshape(-1)
+
+            # sort by expert index so that all objects for each expert are grouped together
+            order = torch.argsort(expert_indices)
+            # update to be in sorted order by expert index
+            object_indices = object_indices[order]
+            expert_indices = expert_indices[order]
+            expert_weights = expert_weights[order]
+            # count how many objects are assigned to each expert to know how to split the input tensor for each expert's forward pass
+            objects_per_expert = torch.bincount(expert_indices, minlength=self.num_experts)
+
+            if not self.training:
+                self.expert_dispatch_counts.add_(objects_per_expert.to(self.expert_dispatch_counts.device))
+
+            cursor = 0
+            # iterate through each expert's assigned objects in order of expert index
+            for expert_id, count in enumerate(objects_per_expert.tolist()):
+                if count == 0:
+                    continue
+                end = cursor + count
+
+                # create minibatch if all objects assigned to the current expert
+                current_object_indices = object_indices[cursor:end]
+                current_inputs = x.index_select(0, current_object_indices)
+
+                # forward pass through the current expert with created minibatch
+                current_outputs = self.routed_experts[expert_id](current_inputs)
+                # get the corresponding expert weights for the current expert's assigned objects
+                current_weights = expert_weights[cursor:end].unsqueeze(-1)
+
+                # weight the expert outputs by the corresponding expert weights for each object,
+                # then add to the correct rows of the final output tensor using the object indices
+                routed_output.index_add_(0, current_object_indices, current_outputs * current_weights)
+
+                cursor = end
+
+        # as shared experts are not part of the routing decisions,
+        # run all objects through all shared experts and add to the final output
+        if self.num_shared_experts > 0:
+            shared_output = torch.zeros_like(routed_output)
+            for shared_expert in self.shared_experts:
+                shared_output = shared_output + shared_expert(x)
+            final_output = routed_output + shared_output
+        else:
+            final_output = routed_output
+
+        # fi - proportion of objects assigned to each expert
+        denom = max(num_objects * self.select_top_k, 1)
+        dispatch_fraction = objects_per_expert.to(dtype=dense_gate_weights.dtype) / denom
+        # pi - average probability of each expert being selected across all objects
+        mean_router_prob = dense_gate_weights.mean(dim=0) if num_objects > 0 else torch.zeros_like(dispatch_fraction)
+        # l_aux is the sum across experts of fi * pi, scaled by alpha and num_experts
+        l_aux = self.alpha * self.num_experts * torch.sum(dispatch_fraction * mean_router_prob)
+
+        if num_objects > 0:
+            # use clean router logits without noise (pre-softmax)
+            cz_lz = self.c_z * torch.mean(torch.logsumexp(router_logits, dim=-1).pow(2))
+        else:
+            cz_lz = torch.zeros((), dtype=x.dtype, device=x.device)
+
+        # convert back to original batch/object shape, with the MoE output in the last dimension
+        final_output = final_output.view(original_shape[0], original_shape[1], self.embed_dim)
+
+        return final_output, l_aux, cz_lz
+
+
 class TransformerBlockModule(nn.Module):
     def __init__(self, projection_dim, num_heads, dropout, talking_head, layer_scale, layer_scale_init,
-                 drop_probability):
+                 drop_probability, use_moe: bool = False, moe_base_num_experts: int = 4,
+                 moe_base_select_top_k: int = 2, moe_num_shared_experts: int = 0,
+                 moe_expert_segmentation_factor: int = 1, moe_scale_expert_dim: bool = False,
+                 moe_alpha: float = 0.01, moe_cz: float = 0.0, moe_use_router_noise: bool = False):
         super().__init__()
         self.projection_dim = projection_dim
         self.num_heads = num_heads
@@ -18,6 +246,7 @@ def __init__(self, projection_dim, num_heads, dropout, talking_head, layer_scale
         self.talking_head = talking_head
         self.layer_scale_flag = layer_scale
         self.drop_probability = drop_probability
+        self.use_moe = use_moe
 
         self.norm1 = nn.LayerNorm(projection_dim)
         self.norm2 = nn.LayerNorm(projection_dim)
@@ -27,13 +256,30 @@ def __init__(self, projection_dim, num_heads, dropout, talking_head, layer_scale
         else:
             self.attn = nn.MultiheadAttention(projection_dim, num_heads, dropout, batch_first=True)
 
-        self.mlp = nn.Sequential(
-            nn.Linear(projection_dim, 2 * projection_dim),
-            nn.GELU(approximate='none'),
-            nn.Dropout(dropout),
-            nn.Linear(2 * projection_dim, projection_dim),
-        )
+        if not self.use_moe:
+            self.mlp = nn.Sequential(
+                nn.Linear(projection_dim, 2 * projection_dim),
+                nn.GELU(approximate="none"),
+                nn.Dropout(dropout),
+                nn.Linear(2 * projection_dim, projection_dim),
+            )
+        else:
+            self.mlp = MoE(
+                embed_dim=projection_dim,
+                feedforward_dim=2 * projection_dim,
+                base_num_experts=moe_base_num_experts,
+                base_select_top_k=moe_base_select_top_k,
+                num_shared_experts=moe_num_shared_experts,
+                expert_segmentation_factor=moe_expert_segmentation_factor,
+                scale_expert_dim=moe_scale_expert_dim,
+                alpha=moe_alpha,
+                c_z=moe_cz,
+                use_router_noise=moe_use_router_noise,
+                dropout=dropout,
+            )
 
+        self.moe_l_aux = torch.tensor(0.0)
+        self.moe_cz_lz = torch.tensor(0.0)
         self.drop_path = StochasticDepth(drop_probability)
 
         if layer_scale:
@@ -42,6 +288,9 @@ def __init__(self, projection_dim, num_heads, dropout, talking_head, layer_scale
 
     def forward(self, x, mask, attn_mask=None):
         # TransformerBlock input shapes: x: torch.Size([B, P, 128]), mask: torch.Size([B, P, 1])
+        self.moe_l_aux = x.new_zeros(())
+        self.moe_cz_lz = x.new_zeros(())
+
         padding_mask = ~(mask.squeeze(2).bool()) if mask is not None else None  # [batch_size, num_objects]
         if self.talking_head:
 
@@ -70,11 +319,19 @@ def forward(self, x, mask, attn_mask=None):
             # Input updates: torch.Size([B, P, 128]), mask: torch.Size([B, P])
             x2 = x + self.drop_path(self.layer_scale1(updates, mask))
             x3 = self.norm2(x2)
-            x = x2 + self.drop_path(self.layer_scale2(self.mlp(x3), mask))
+            if self.use_moe:
+                x4, self.moe_l_aux, self.moe_cz_lz = self.mlp(x3)
+            else:
+                x4 = self.mlp(x3)
+            x = x2 + self.drop_path(self.layer_scale2(x4, mask))
         else:
             x2 = x + self.drop_path(updates)
             x3 = self.norm2(x2)
-            x = x2 + self.drop_path(self.mlp(x3))
+            if self.use_moe:
+                x4, self.moe_l_aux, self.moe_cz_lz = self.mlp(x3)
+            else:
+                x4 = self.mlp(x3)
+            x = x2 + self.drop_path(x4)
 
         if mask is not None:
             x = x * mask
@@ -396,3 +653,60 @@ def forward(self,
         tgt = self.norm3(tgt)
         return tgt
 
+
+def log_moe_expert_distribution(
+    model: nn.Module,
+    reset: bool = True,
+    logger: Optional[logging.Logger] = None,
+) -> None:
+    """Log the expert dispatch distribution for every MoE layer in *model*.
+
+    Call this after running eval batches to see how evenly inputs were routed.
+    Counts are accumulated across all eval forward passes since the last reset.
+
+    Args:
+        model:  Any nn.Module that may contain MoE sub-modules.
+        reset:  If True (default), zero the counts after logging so the next
+                eval epoch starts fresh.
+        logger: Logger to write to.  Defaults to this module's logger.
+
+    Example usage in an eval loop::
+
+        model.eval()
+        for batch in eval_loader:
+            with torch.no_grad():
+                model(batch)
+        log_moe_expert_distribution(model)   # prints distribution, resets counts
+    """
+    _log = (logger.info if logger is not None else print)
+    _warn = (logger.warning if logger is not None else print)
+
+    moe_layers = [(name, m) for name, m in model.named_modules() if isinstance(m, MoE)]
+    if not moe_layers:
+        _warn("log_moe_expert_distribution: no MoE layers found in model.")
+        return
+
+    for name, moe in moe_layers:
+        counts = moe.expert_dispatch_counts.cpu()
+        total = counts.sum().item()
+
+        if total == 0:
+            _log(f"[MoE dist] {name}: no eval data recorded (counts are all zero).")
+            if reset:
+                moe.reset_expert_dispatch_counts()
+            continue
+
+        uniform = total / moe.num_experts
+        lines = [
+            f"[MoE dist] {name}  "
+            f"(total dispatches={total:,}, ideal per expert={uniform:,.1f})"
+        ]
+        for i, c in enumerate(counts.tolist()):
+            pct = 100.0 * c / total
+            bar = "█" * int(pct / 2)  # each block ≈ 2 %
+            lines.append(f"  expert {i:3d}: {c:9,d}  ({pct:5.1f}%)  {bar}")
+        _log("\n".join(lines))
+
+        if reset:
+            moe.reset_expert_dispatch_counts()
+
diff --git a/network/loss/moe.py b/network/loss/moe.py
new file mode 100644
index 0000000..cfc8988
--- /dev/null
+++ b/network/loss/moe.py
@@ -0,0 +1,25 @@
+from torch import Tensor
+from typing import Optional, Tuple, Dict, Any
+
+
+def loss(model_output: Dict[str, Any]) -> Tuple[Optional[Tensor], Optional[Tensor]]:
+    """
+    Extract MoE auxiliary losses from model output.
+
+    The auxiliary load-balancing loss (l_aux) and z-loss (cz_lz) are computed
+    during the model forward pass inside each MoE transformer block and
+    accumulated onto the model output dict under the keys "L_aux" and "cz_Lz".
+    Callers should add both directly to their total loss without additional
+    scaling; the weights are already baked in via the alpha and c_z hyperparameters
+    configured on each MoE layer.
+
+    Args:
+        model_output: dict returned by EveNetModel.forward()
+
+    Returns:
+        l_aux:  auxiliary load-balancing loss, or None when MoE is disabled
+        cz_lz:  z-loss regularisation term, or None when MoE is disabled
+    """
+    l_aux = model_output.get("L_aux", None)
+    cz_lz = model_output.get("cz_Lz", None)
+    return l_aux, cz_lz
diff --git a/utilities/logger.py b/utilities/logger.py
index 4c5ba3e..d9846b3 100644
--- a/utilities/logger.py
+++ b/utilities/logger.py
@@ -123,6 +123,6 @@ def setup_logging(log_level=logging.INFO, rank: int = 0, log_dir="logs"):
         format=f"%(asctime)s | %(name)s | %(levelname)s | %(message)s",
         handlers=[
             logging.FileHandler(log_file, mode="a"),
-            # logging.StreamHandler() if rank == 0 else logging.NullHandler(),  # Only rank 0 logs to stdout
+            logging.StreamHandler() if rank == 0 else logging.NullHandler(),
         ]
     )