diff --git a/egs/libriheavy/LM/zipformer1/chunk_decoder.py b/egs/libriheavy/LM/zipformer1/chunk_decoder.py
new file mode 100644
index 000000000..5f3a84be0
--- /dev/null
+++ b/egs/libriheavy/LM/zipformer1/chunk_decoder.py
@@ -0,0 +1,117 @@
+#!/usr/bin/env python3
+# Copyright    2023       Xiaomi Corp.        (authors: Daniel Povey)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+
+import torch
+from torch import nn, Tensor
+
+
+class ChunkDecoder(nn.Module):
+    """
+    """
+    def __init__(self,
+                 embed_dim: int,
+                 chunk_size: int,
+                 vocab_size: int,
+                 hidden_size: int,
+                 num_layers: int = 2):
+        """
+        A 'decoder' that computes the probability of symbols in a language modeling task.
+        Conceptually it computes the probability of `chunk_size` symbols (e.g. 8 symbols)
+        based on an embedding derived from all symbols preceding this chunk of 8 symbols.
+        Also, within the chunk, we always see all previous symbols (plus the last symbol
+        of the previous chunk).
+        """
+        super().__init__()
+        self.chunk_size = chunk_size
+        self.num_layers = num_layers
+        self.hidden_size = hidden_size
+
+        self.lstm = nn.LSTM(input_size=embed_dim,
+                            hidden_size=hidden_size,
+                            num_layers=num_layers)
+
+        self.label_embed = nn.Embedding(
+            num_embeddings=vocab_size,
+            embedding_dim=embed_dim)
+
+        # project to hidden state and cell state at the beginning of each chunk.
+        # (we don't run the lstm contiuously over the sequence, for both
+        # training speed and stability; instead, we reconstruct the hidden
+        # state for each chunk.)
+        # the factor of 2 is to cover hidden state and cell state.
+        self.init_proj = nn.Linear(embed_dim,
+                                   2 * hidden_size * num_layers)
+
+        self.out_proj = nn.Linear(hidden_size,
+                                  vocab_size)
+
+
+    def forward(self,
+                labels: Tensor,
+                encoder_embed: Tensor) -> Tensor:
+        """
+        Compute log-probs.
+        Args:
+           labels: the labels, a Tensor of integer type of shape (batch_size, seq_len);
+                 seq_len is expected to be a multiple of chunk_size.
+        encoder_embed: the embeddings from the encoder, of shape (seq_len//chunk_size, batch_size, embed_dim)
+
+        Returns:
+            returns the log-probs for each symbol, in a Tensor of shape (batch_size, seq_len).
+        """
+        (batch_size, seq_len) = labels.shape
+        (num_chunks, _batch_size, embed_dim) = encoder_embed.shape
+        chunk_size = self.chunk_size
+        assert batch_size == _batch_size
+        assert num_chunks * chunk_size == seq_len
+
+        labels_shifted = torch.cat((torch.zeros_like(labels[0:1]),
+                                    labels[:-1]), dim=0)
+
+        labels_embed = self.label_embed(labels_shifted.t())  # (seq_len, batch_size, embed_dim)
+
+        init = self.init_proj(encoder_embed).reshape(num_chunks, batch_size,
+                                                     2, self.num_layers, self.hidden_size)
+        init = init.permute(2, 3, 0, 1, 4).reshape(2, self.num_layers,
+                                                   num_chunks * batch_size,
+                                                   self.hidden_size).contiguous()
+        hidden = init[0]
+        cell = init[1]
+
+
+        labels_embed = labels_embed.reshape(num_chunks, chunk_size, batch_size, embed_dim).transpose(0, 1)
+        labels_embed = labels_embed.contiguous().reshape(chunk_size, num_chunks * batch_size, embed_dim)
+        encoder_embed = encoder_embed.reshape(1, num_chunks * batch_size, embed_dim)
+
+        x = labels_embed + encoder_embed  # broadcasts encoder_embed over the chunk_size
+
+        (x, _hidden) = self.lstm(x, hx=(hidden, cell))
+
+        x = self.out_proj(x)
+
+        vocab_size = x.shape[-1]
+        # x: (chunk_size, num_chunks * batch_size, vocab_size)
+        x = x.reshape(chunk_size, num_chunks, batch_size, vocab_size)
+        x = x.permute(2, 1, 0, 3).contiguous().reshape(batch_size, num_chunks * chunk_size, vocab_size)
+
+        x = x.log_softmax(dim=-1)
+
+        logprobs = torch.gather(x, dim=-1, index=labels.unsqueeze(-1)).squeeze(-1)  # (batch_size, seq_len)
+
+        return logprobs
diff --git a/egs/libriheavy/LM/zipformer1/encoder_interface.py b/egs/libriheavy/LM/zipformer1/encoder_interface.py
new file mode 120000
index 000000000..dbaf9d4a8
--- /dev/null
+++ b/egs/libriheavy/LM/zipformer1/encoder_interface.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless/encoder_interface.py
\ No newline at end of file
diff --git a/egs/libriheavy/LM/zipformer1/lm_datamodule.py b/egs/libriheavy/LM/zipformer1/lm_datamodule.py
new file mode 100644
index 000000000..8a269a179
--- /dev/null
+++ b/egs/libriheavy/LM/zipformer1/lm_datamodule.py
@@ -0,0 +1,141 @@
+# Copyright      2021  Piotr Żelasko
+# Copyright      2022  Xiaomi Corporation     (Author: Mingshuang Luo)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import argparse
+import inspect
+import logging
+from functools import lru_cache
+import numpy as np
+import random
+from pathlib import Path
+from typing import Any, Dict, Optional
+from icefall.dist import get_world_size, get_rank
+
+import torch
+
+from torch.utils.data import DataLoader
+
+from icefall.utils import str2bool
+
+
+
+class LmDataset(torch.utils.data.IterableDataset):
+    def __init__(self,
+                 file_list_fn: Path,
+                 bytes_per_segment: int = 200):
+        """
+        Initialize LmDataset object.  Args:
+          file_list_fn: a file in which each line contains: a number of bytes, then a space, then a filename.
+              e.g. a line might contain the text "64324 foo/abc.txt".
+              (filenames can not contain spaces).
+          bytes_per_segment: the number of bytes in each segment of data.
+        """
+        self.files = []
+        self.num_bytes = []
+        self.bytes_per_segment = bytes_per_segment
+
+        num_bytes = []
+        with open(file_list_fn) as f:
+            for line in f.readlines():
+                line = line.strip()  # remove newline
+                num_bytes = line.split()[0]  # a str
+                fn = line[len(num_bytes) + 1:]  # this works even if fn has spaces in
+                self.files.append(fn)
+                self.num_bytes.append(int(num_bytes))
+        tot_bytes = sum(self.num_bytes)
+        N = len(self.num_bytes)
+        self.probs = np.array([ x / tot_bytes for x in self.num_bytes ])
+
+        worker_info = torch.utils.data.get_worker_info()
+        num_workers = (1 if worker_info is None else worker_info.num_workers)
+
+        tot_workers = num_workers * get_world_size()
+
+        self.num_segments = tot_bytes // (bytes_per_segment * tot_workers)
+
+
+    def __iter__(self):
+        worker_info = torch.utils.data.get_worker_info()
+        # id includes both worker (within training job) and rank of training job
+        my_id = (0 if worker_info is None else worker_info.id) + 1000 * get_rank()
+
+        seed = random.randint(0, 10000) + my_id
+        logging.info(f"seed={seed}, num_segments={self.num_segments}")
+        rng = np.random.default_rng(seed=seed)
+        for n in range(self.num_segments):
+            # np.random.multinomial / np.random.Generator.multinomial has an interface
+            # where it gives counts of different categories, instead of the chosen category,
+            # so we need to use np.nonzero to get the chosen category (i.e. the file index)
+            # np.nonzero will give an array per dim, so file_idx,
+            # gives the array of nonzero index
+            file_idx, = np.nonzero(rng.multinomial(1, self.probs))
+            file_idx, = file_idx
+
+            fn = self.files[file_idx]
+            num_bytes = self.num_bytes[file_idx]
+
+            # begin_pos, end_pos are the begin,end of a range from which we'll pick
+            # randomly, for where the start of the segment might be.
+            begin_pos = 0
+            end_pos = max(1, num_bytes - self.bytes_per_segment)
+
+            begin, = rng.integers(low=begin_pos, high=end_pos, size=1)
+
+            with open(fn, "rb") as f:
+                f.seek(begin)
+                b = f.read(self.bytes_per_segment) # b is bytes object
+            read_size = len(b)
+            if read_size < self.bytes_per_segment:
+                b = b + b'\0' * (self.bytes_per_segment - read_size)
+            yield torch.Tensor(np.frombuffer(b, dtype=np.uint8).copy()).to(torch.long)
+
+
+
+def LmDataloader(dataset: LmDataset,
+                 batch_size: int,
+                 num_workers: int):
+
+    return torch.utils.data.DataLoader(
+        dataset=dataset,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        drop_last=True)
+
+
+
+
+def _test():
+    l = LmDataset('files.txt')
+
+    d = LmDataloader(l, batch_size=5, num_workers=4)
+
+    for batch in d:
+        logging.info("batch shape: ", batch.shape)
+
+
+
+if __name__ == '__main__':
+    logging.getLogger().setLevel(logging.INFO)
+    _test()
+
+
+
+# cd libriheavy/LM
+# find /ceph-data3/xiaoyu/librilight_text/output_text_large_cleaned -name text.txt -exec stat --printf='%s ' {} \; -print > files.txt
+# head -n 2 files.txt > valid.txt
+# tail -n +3 files.txt > train.txt
diff --git a/egs/libriheavy/LM/zipformer1/model.py b/egs/libriheavy/LM/zipformer1/model.py
new file mode 100644
index 000000000..43fc715dd
--- /dev/null
+++ b/egs/libriheavy/LM/zipformer1/model.py
@@ -0,0 +1,65 @@
+#!/usr/bin/env python3
+# Copyright    2023       Xiaomi Corp.        (authors: Daniel Povey)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+
+import torch
+from torch import nn, Tensor
+from chunk_decoder import ChunkDecoder
+from zipformer import Zipformer2
+
+
+class Zipformer2LM(nn.Module):
+
+    def __init__(self,
+                 encoder_embed: nn.Module,
+                 encoder: Zipformer2,
+                 decoder: ChunkDecoder):
+        super().__init__()
+        self.encoder_embed = encoder_embed
+        self.encoder = encoder # does subsampling
+        self.decoder = decoder
+
+
+    def forward(self,
+                labels: Tensor):
+        """
+        Compute array of log-probs
+
+        Args:
+         labels: a Tensor containing the labels (in the range 0..num_symbols-1), of shape (batch_size, seq_len).
+        Returns:
+           a Tensor containing the log-probs for each label, of shape (batch_size, seq_len).
+        """
+        (batch_size, seq_len) = labels.shape
+
+        chunk_size = self.decoder.chunk_size
+        labels_shifted = labels.t()  # (time, batch)
+        labels_shifted = torch.cat((torch.zeros_like(labels_shifted[:chunk_size]),
+                                    labels_shifted[:-chunk_size]),
+                                   dim=0)
+
+        x = self.encoder_embed(labels_shifted)
+        x_lens = torch.full((batch_size,), seq_len,
+                            dtype=torch.long, device=labels.device)
+        # x_lens is after subsampling.  Actually we don't need it.
+
+
+        (x, x_lens) = self.encoder(x, x_lens)
+
+        logprobs = self.decoder(labels, x)
+        return logprobs
diff --git a/egs/libriheavy/LM/zipformer1/optim.py b/egs/libriheavy/LM/zipformer1/optim.py
new file mode 120000
index 000000000..43a249f3a
--- /dev/null
+++ b/egs/libriheavy/LM/zipformer1/optim.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/zipformer2/optim.py
\ No newline at end of file
diff --git a/egs/libriheavy/LM/zipformer1/scaling.py b/egs/libriheavy/LM/zipformer1/scaling.py
new file mode 120000
index 000000000..1158af7c8
--- /dev/null
+++ b/egs/libriheavy/LM/zipformer1/scaling.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/zipformer2/scaling.py
\ No newline at end of file
diff --git a/egs/libriheavy/LM/zipformer1/train.py b/egs/libriheavy/LM/zipformer1/train.py
new file mode 100755
index 000000000..5f0715be6
--- /dev/null
+++ b/egs/libriheavy/LM/zipformer1/train.py
@@ -0,0 +1,1178 @@
+#!/usr/bin/env python3
+# Copyright    2021-2022  Xiaomi Corp.        (authors: Fangjun Kuang,
+#                                                       Wei Kang,
+#                                                       Mingshuang Luo,)
+#                                                       Zengwei Yao)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Usage:
+
+export CUDA_VISIBLE_DEVICES="0,1,2,3"
+
+./pruned_transducer_stateless7/train.py \
+  --world-size 4 \
+  --num-epochs 30 \
+  --start-epoch 1 \
+  --exp-dir pruned_transducer_stateless7/exp \
+  --full-libri 1 \
+  --max-duration 300
+
+# For mix precision training:
+
+./pruned_transducer_stateless7/train.py \
+  --world-size 4 \
+  --num-epochs 30 \
+  --start-epoch 1 \
+  --use-fp16 1 \
+  --exp-dir pruned_transducer_stateless7/exp \
+  --full-libri 1 \
+  --max-duration 550
+
+"""
+
+
+import argparse
+import copy
+import logging
+import random
+import warnings
+from pathlib import Path
+from shutil import copyfile
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import k2
+import numpy
+import optim
+import torch
+import torch.multiprocessing as mp
+import torch.nn as nn
+from lm_datamodule import LmDataset, LmDataloader
+from zipformer import Zipformer2
+from scaling import ScheduledFloat
+from lhotse.utils import fix_random_seed
+from chunk_decoder import ChunkDecoder
+from model import Zipformer2LM
+from optim import Eden, ScaledAdam
+from torch import Tensor
+from torch import nn
+from torch.cuda.amp import GradScaler
+
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.tensorboard import SummaryWriter
+
+from icefall import diagnostics
+from icefall.checkpoint import load_checkpoint, remove_checkpoints
+from icefall.checkpoint import save_checkpoint as save_checkpoint_impl
+from icefall.checkpoint import (
+    save_checkpoint_with_global_batch_idx,
+    update_averaged_model,
+)
+from icefall.hooks import register_inf_check_hooks
+from icefall.dist import cleanup_dist, setup_dist
+from icefall.env import get_env_info
+from icefall.utils import (
+    AttributeDict,
+    MetricsTracker,
+    setup_logger,
+    str2bool,
+    get_parameter_groups_with_lrs
+)
+
+LRSchedulerType = Union[
+    torch.optim.lr_scheduler._LRScheduler, optim.LRScheduler
+]
+
+
+def get_adjusted_batch_count(
+        params: AttributeDict) -> float:
+    # don't do any adjustment for now.
+    # This is for purposes of set_batch_count().
+    return params.batch_idx_train
+
+
+
+def set_batch_count(
+        model: Union[nn.Module, DDP], batch_count: float
+) -> None:
+    if isinstance(model, DDP):
+        # get underlying nn.Module
+        model = model.module
+    for name, module in model.named_modules():
+        if hasattr(module, 'batch_count'):
+            module.batch_count = batch_count
+        if hasattr(module, 'name'):
+            module.name = name
+
+
+def add_model_arguments(parser: argparse.ArgumentParser):
+    parser.add_argument(
+        "--num-encoder-layers",
+        type=str,
+        default="2,4,5,6",
+        help="Number of zipformer encoder layers per stack, comma separated.",
+    )
+
+
+    parser.add_argument(
+        "--downsampling-factor",
+        type=str,
+        default="1,2,4,8",
+        help="Downsampling factor for each stack of encoder layers.",
+    )
+
+
+    parser.add_argument(
+        "--feedforward-dim",
+        type=str,
+        default="512,768,1024,1536",
+        help="Feedforward dimension of the zipformer encoder layers, per stack, comma separated.",
+    )
+
+    parser.add_argument(
+        "--num-heads",
+        type=str,
+        default="4,4,6,8",
+        help="Number of attention heads in the zipformer encoder layers: a single int or comma-separated list.",
+    )
+
+    parser.add_argument(
+        "--encoder-dim",
+        type=str,
+        default="192,256,384,512",
+        help="Embedding dimension in encoder stacks: a single int or comma-separated list."
+    )
+
+    parser.add_argument(
+        "--query-head-dim",
+        type=str,
+        default="32",
+        help="Query/key dimension per head in encoder stacks: a single int or comma-separated list."
+    )
+
+    parser.add_argument(
+        "--value-head-dim",
+        type=str,
+        default="12",
+        help="Value dimension per head in encoder stacks: a single int or comma-separated list."
+    )
+
+    parser.add_argument(
+        "--pos-head-dim",
+        type=str,
+        default="4",
+        help="Positional-encoding dimension per head in encoder stacks: a single int or comma-separated list."
+    )
+
+    parser.add_argument(
+        "--pos-dim",
+        type=int,
+        default="48",
+        help="Positional-encoding embedding dimension"
+    )
+
+    parser.add_argument(
+        "--encoder-unmasked-dim",
+        type=str,
+        default="192,192,256,256",
+        help="Unmasked dimensions in the encoders, relates to augmentation during training.  "
+        "A single int or comma-separated list.  Must be <= each corresponding encoder_dim."
+    )
+
+    parser.add_argument(
+        "--cnn-module-kernel",
+        type=str,
+        default="31,31,15,15",
+        help="Sizes of convolutional kernels in convolution modules in each encoder stack: "
+        "a single int or comma-separated list.",
+    )
+
+    parser.add_argument(
+        "--decoder-hidden-size",
+        type=int,
+        default=768,
+        help="LSTM dimension  in decoder",
+    )
+
+    parser.add_argument(
+        "--decoder-num-layers",
+        type=int,
+        default=2,
+        help="Number of LSTM layers in decoder",
+    )
+
+
+
+def get_parser():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+
+    parser.add_argument(
+        "--world-size",
+        type=int,
+        default=1,
+        help="Number of GPUs for DDP training.",
+    )
+
+    parser.add_argument(
+        "--master-port",
+        type=int,
+        default=12354,
+        help="Master port to use for DDP training.",
+    )
+
+    parser.add_argument(
+        "--tensorboard",
+        type=str2bool,
+        default=True,
+        help="Should various information be logged in tensorboard.",
+    )
+
+    parser.add_argument(
+        "--num-epochs",
+        type=int,
+        default=30,
+        help="Number of epochs to train.",
+    )
+
+    parser.add_argument(
+        "--start-epoch",
+        type=int,
+        default=1,
+        help="""Resume training from this epoch. It should be positive.
+        If larger than 1, it will load checkpoint from
+        exp-dir/epoch-{start_epoch-1}.pt
+        """,
+    )
+
+    parser.add_argument(
+        "--start-batch",
+        type=int,
+        default=0,
+        help="""If positive, --start-epoch is ignored and
+        it loads the checkpoint from exp-dir/checkpoint-{start_batch}.pt
+        """,
+    )
+
+    parser.add_argument(
+        "--exp-dir",
+        type=str,
+        default="pruned_transducer_stateless7/exp",
+        help="""The experiment dir.
+        It specifies the directory where all training related
+        files, e.g., checkpoints, log, etc, are saved
+        """,
+    )
+
+
+
+    parser.add_argument(
+        "--base-lr",
+        type=float,
+        default=0.045,
+        help="The base learning rate."
+    )
+
+    parser.add_argument(
+        "--lr-batches",
+        type=float,
+        default=7500,
+        help="""Number of steps that affects how rapidly the learning rate
+        decreases. We suggest not to change this.""",
+    )
+
+    parser.add_argument(
+        "--lr-epochs",
+        type=float,
+        default=3.5,
+        help="""Number of epochs that affects how rapidly the learning rate decreases.
+        """,
+    )
+
+    parser.add_argument(
+        "--ref-duration",
+        type=float,
+        default=600,
+        help="Reference batch duration for purposes of adjusting batch counts for setting various "
+        "schedules inside the model"
+    )
+
+
+    parser.add_argument(
+        "--context-size",
+        type=int,
+        default=2,
+        help="The context size in the decoder. 1 means bigram; "
+        "2 means tri-gram",
+    )
+
+    parser.add_argument(
+        "--prune-range",
+        type=int,
+        default=5,
+        help="The prune range for rnnt loss, it means how many symbols(context)"
+        "we are using to compute the loss",
+    )
+
+    parser.add_argument(
+        "--lm-scale",
+        type=float,
+        default=0.25,
+        help="The scale to smooth the loss with lm "
+        "(output of prediction network) part.",
+    )
+
+    parser.add_argument(
+        "--am-scale",
+        type=float,
+        default=0.0,
+        help="The scale to smooth the loss with am (output of encoder network)"
+        "part.",
+    )
+
+    parser.add_argument(
+        "--simple-loss-scale",
+        type=float,
+        default=0.5,
+        help="To get pruning ranges, we will calculate a simple version"
+        "loss(joiner is just addition), this simple loss also uses for"
+        "training (as a regularization item). We will scale the simple loss"
+        "with this parameter before adding to the final loss.",
+    )
+
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="The seed for random generators intended for reproducibility",
+    )
+
+    parser.add_argument(
+        "--print-diagnostics",
+        type=str2bool,
+        default=False,
+        help="Accumulate stats on activations, print them and exit.",
+    )
+
+    parser.add_argument(
+        "--inf-check",
+        type=str2bool,
+        default=False,
+        help="Add hooks to check for infinite module outputs and gradients.",
+    )
+
+    parser.add_argument(
+        "--save-every-n",
+        type=int,
+        default=4000,
+        help="""Save checkpoint after processing this number of batches"
+        periodically. We save checkpoint to exp-dir/ whenever
+        params.batch_idx_train % save_every_n == 0. The checkpoint filename
+        has the form: f'exp-dir/checkpoint-{params.batch_idx_train}.pt'
+        Note: It also saves checkpoint to `exp-dir/epoch-xxx.pt` at the
+        end of each epoch where `xxx` is the epoch number counting from 0.
+        """,
+    )
+
+    parser.add_argument(
+        "--keep-last-k",
+        type=int,
+        default=30,
+        help="""Only keep this number of checkpoints on disk.
+        For instance, if it is 3, there are only 3 checkpoints
+        in the exp-dir with filenames `checkpoint-xxx.pt`.
+        It does not affect checkpoints with name `epoch-xxx.pt`.
+        """,
+    )
+
+    parser.add_argument(
+        "--average-period",
+        type=int,
+        default=200,
+        help="""Update the averaged model, namely `model_avg`, after processing
+        this number of batches. `model_avg` is a separate version of model,
+        in which each floating-point parameter is the average of all the
+        parameters from the start of training. Each time we take the average,
+        we do: `model_avg = model * (average_period / batch_idx_train) +
+            model_avg * ((batch_idx_train - average_period) / batch_idx_train)`.
+        """,
+    )
+
+    parser.add_argument(
+        "--use-fp16",
+        type=str2bool,
+        default=False,
+        help="Whether to use half precision training.",
+    )
+
+    add_model_arguments(parser)
+
+    return parser
+
+
+def get_params() -> AttributeDict:
+    """Return a dict containing training parameters.
+
+    All training related parameters that are not passed from the commandline
+    are saved in the variable `params`.
+
+    Commandline options are merged into `params` after they are parsed, so
+    you can also access them via `params`.
+
+    Explanation of options saved in `params`:
+
+        - best_train_loss: Best training loss so far. It is used to select
+                           the model that has the lowest training loss. It is
+                           updated during the training.
+
+        - best_valid_loss: Best validation loss so far. It is used to select
+                           the model that has the lowest validation loss. It is
+                           updated during the training.
+
+        - best_train_epoch: It is the epoch that has the best training loss.
+
+        - best_valid_epoch: It is the epoch that has the best validation loss.
+
+        - batch_idx_train: Used to writing statistics to tensorboard. It
+                           contains number of batches trained so far across
+                           epochs.
+
+        - log_interval:  Print training loss if batch_idx % log_interval` is 0
+
+        - reset_interval: Reset statistics if batch_idx % reset_interval is 0
+
+        - valid_interval:  Run validation if batch_idx % valid_interval is 0
+
+        - feature_dim: The model input dim. It has to match the one used
+                       in computing features.
+
+        - subsampling_factor:  The subsampling factor for the model.
+
+        - encoder_dim: Hidden dim for multi-head attention model.
+
+        - num_decoder_layers: Number of decoder layer of transformer decoder.
+
+        - warm_step: The warmup period that dictates the decay of the
+              scale on "simple" (un-pruned) loss.
+    """
+    params = AttributeDict(
+        {
+            "best_train_loss": float("inf"),
+            "best_valid_loss": float("inf"),
+            "best_train_epoch": -1,
+            "best_valid_epoch": -1,
+            "batch_idx_train": 0,
+            "log_interval": 50,
+            "reset_interval": 200,
+            "valid_interval": 3000,
+            "warm_step": 2000,
+            "env_info": get_env_info(),
+            "bytes_per_segment": 1024,
+            "batch_size": 64,
+            "train_file_list": "train.txt",
+            "valid_file_list": "valid.txt",
+            "num_workers": 4,
+        }
+    )
+
+    return params
+
+
+def _to_int_tuple(s: str):
+    return tuple(map(int, s.split(',')))
+
+
+
+def get_encoder_embed(params: AttributeDict) -> nn.Module:
+    return nn.Embedding(
+        num_embeddings=256,   # we encode the text as UTF-8 bytes
+        embedding_dim=_to_int_tuple(params.encoder_dim)[0],
+    )
+
+
+
+def get_encoder_model(params: AttributeDict) -> nn.Module:
+    chunk_size = _to_int_tuple(params.downsampling_factor)[-1]
+    encoder = Zipformer2(
+        output_downsampling_factor=chunk_size,
+        downsampling_factor=_to_int_tuple(params.downsampling_factor),
+        num_encoder_layers=_to_int_tuple(params.num_encoder_layers),
+        encoder_dim=_to_int_tuple(params.encoder_dim),
+        encoder_unmasked_dim=_to_int_tuple(params.encoder_unmasked_dim),
+        query_head_dim=_to_int_tuple(params.query_head_dim),
+        pos_head_dim=_to_int_tuple(params.pos_head_dim),
+        value_head_dim=_to_int_tuple(params.value_head_dim),
+        pos_dim=params.pos_dim,
+        num_heads=_to_int_tuple(params.num_heads),
+        feedforward_dim=_to_int_tuple(params.feedforward_dim),
+        cnn_module_kernel=_to_int_tuple(params.cnn_module_kernel),
+        dropout=ScheduledFloat((0.0, 0.3), (20000.0, 0.1)),
+        warmup_batches=4000.0,
+        causal=True,
+        chunk_size=(chunk_size,),
+        left_context_frames=(-1,),
+    )
+    return encoder
+
+
+def get_decoder_model(params: AttributeDict) -> nn.Module:
+    chunk_size = _to_int_tuple(params.downsampling_factor)[-1]
+    decoder = ChunkDecoder(
+        embed_dim=max(_to_int_tuple(params.encoder_dim)),
+        chunk_size=chunk_size,
+        vocab_size=256, # bytes
+        hidden_size=params.decoder_hidden_size,
+        num_layers=params.decoder_num_layers,
+    )
+    return decoder
+
+
+def get_model(params: AttributeDict) -> nn.Module:
+    encoder_embed = get_encoder_embed(params)
+    encoder = get_encoder_model(params)
+    decoder = get_decoder_model(params)
+
+    model = Zipformer2LM(
+        encoder_embed=encoder_embed,
+        encoder=encoder,
+        decoder=decoder,
+    )
+    return model
+
+
+def load_checkpoint_if_available(
+    params: AttributeDict,
+    model: nn.Module,
+    model_avg: nn.Module = None,
+    optimizer: Optional[torch.optim.Optimizer] = None,
+    scheduler: Optional[LRSchedulerType] = None,
+) -> Optional[Dict[str, Any]]:
+    """Load checkpoint from file.
+
+    If params.start_batch is positive, it will load the checkpoint from
+    `params.exp_dir/checkpoint-{params.start_batch}.pt`. Otherwise, if
+    params.start_epoch is larger than 1, it will load the checkpoint from
+    `params.start_epoch - 1`.
+
+    Apart from loading state dict for `model` and `optimizer` it also updates
+    `best_train_epoch`, `best_train_loss`, `best_valid_epoch`,
+    and `best_valid_loss` in `params`.
+
+    Args:
+      params:
+        The return value of :func:`get_params`.
+      model:
+        The training model.
+      model_avg:
+        The stored model averaged from the start of training.
+      optimizer:
+        The optimizer that we are using.
+      scheduler:
+        The scheduler that we are using.
+    Returns:
+      Return a dict containing previously saved training info.
+    """
+    if params.start_batch > 0:
+        filename = params.exp_dir / f"checkpoint-{params.start_batch}.pt"
+    elif params.start_epoch > 1:
+        filename = params.exp_dir / f"epoch-{params.start_epoch-1}.pt"
+    else:
+        return None
+
+    assert filename.is_file(), f"{filename} does not exist!"
+
+    saved_params = load_checkpoint(
+        filename,
+        model=model,
+        model_avg=model_avg,
+        optimizer=optimizer,
+        scheduler=scheduler,
+    )
+
+    keys = [
+        "best_train_epoch",
+        "best_valid_epoch",
+        "batch_idx_train",
+        "best_train_loss",
+        "best_valid_loss",
+    ]
+    for k in keys:
+        params[k] = saved_params[k]
+
+    if params.start_batch > 0:
+        if "cur_epoch" in saved_params:
+            params["start_epoch"] = saved_params["cur_epoch"]
+
+        if "cur_batch_idx" in saved_params:
+            params["cur_batch_idx"] = saved_params["cur_batch_idx"]
+
+    return saved_params
+
+
+def save_checkpoint(
+    params: AttributeDict,
+    model: Union[nn.Module, DDP],
+    model_avg: Optional[nn.Module] = None,
+    optimizer: Optional[torch.optim.Optimizer] = None,
+    scheduler: Optional[LRSchedulerType] = None,
+    scaler: Optional[GradScaler] = None,
+    rank: int = 0,
+) -> None:
+    """Save model, optimizer, scheduler and training stats to file.
+
+    Args:
+      params:
+        It is returned by :func:`get_params`.
+      model:
+        The training model.
+      model_avg:
+        The stored model averaged from the start of training.
+      optimizer:
+        The optimizer used in the training.
+      scaler:
+        The scaler used for mix precision training.
+    """
+    if rank != 0:
+        return
+    filename = params.exp_dir / f"epoch-{params.cur_epoch}.pt"
+    save_checkpoint_impl(
+        filename=filename,
+        model=model,
+        model_avg=model_avg,
+        params=params,
+        optimizer=optimizer,
+        scheduler=scheduler,
+        scaler=scaler,
+        rank=rank,
+    )
+
+    if params.best_train_epoch == params.cur_epoch:
+        best_train_filename = params.exp_dir / "best-train-loss.pt"
+        copyfile(src=filename, dst=best_train_filename)
+
+    if params.best_valid_epoch == params.cur_epoch:
+        best_valid_filename = params.exp_dir / "best-valid-loss.pt"
+        copyfile(src=filename, dst=best_valid_filename)
+
+def _encode_texts_as_bytes(texts: List[str], device: torch.device) -> Tuple[Tensor, Tensor, Tensor]:
+    """
+    Encode texts as bytes and then integer tensors.
+    Args:
+          texts: the texts to encode, as a list of strings
+         device: the PyTorch device we want the texts on
+ Returns:
+       (text, text_lens, style_lens), where:
+    text: a torch.Tensor of shape (batch_size, text_len) containing integers
+       0 <= i < 256
+    text_lens: a torch.Tensor of shape (batch_size,), giving the length of each byt
+       sequence
+    style_lens: a torch.Tensor of shape (batch_size,), giving the length of each
+       style prompt (style prompts are supposed to come first).  Since there is no
+       style prompt here, this is just all zeros.
+    """
+    texts = [ bytes(s, 'UTF-8') for s in texts ]
+    N = len(texts)
+    lengths = [ len(s) for s in texts ]
+    max_len = max(lengths)
+    texts = [ s + (b'\0' * (max_len - len(s))) for s in texts ]
+    text = b''.join(texts)  # bytes array containing all of the texts
+
+    text = torch.Tensor(numpy.frombuffer(text, dtype=numpy.uint8)).to(device)
+    text = text.to(dtype=torch.long)
+    text = text.reshape(N, max_len)
+    text_lens = torch.tensor(lengths).to(device)
+    style_lens = torch.zeros(N, dtype=torch.long, device=device)
+    # print(f"text={text}, text_lens={text_lens}, style_lens={style_lens}")
+    return text, text_lens, style_lens
+
+def compute_loss(
+    params: AttributeDict,
+    model: Union[nn.Module, DDP],
+    batch: Tensor,
+    is_training: bool,
+) -> Tuple[Tensor, MetricsTracker]:
+    """
+    Compute cross-entropy loss given the model and its inputs.
+
+    Args:
+      params:
+        Parameters for training. See :func:`get_params`.
+      model:
+        The model for training. It is an instance of Zipformer in our case.
+      batch:
+        A batch of data. See `lhotse.dataset.K2SpeechRecognitionDataset()`
+        for the content in it.
+      is_training:
+        True for training. False for validation. When it is True, this
+        function enables autograd during computation; when it is False, it
+        disables autograd.
+     warmup: a floating point value which increases throughout training;
+        values >= 1.0 are fully warmed up and have all modules present.
+    """
+    device = (
+        model.device
+        if isinstance(model, DDP)
+        else next(model.parameters()).device
+    )
+
+    labels = batch.to(device)  # (batch_size, sequence_length)
+
+    with torch.set_grad_enabled(is_training):
+        loglikes = model(labels)
+
+        loss = -loglikes.sum()
+
+
+    assert loss.requires_grad == is_training
+
+    info = MetricsTracker()
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        info["frames"] = (
+            labels.numel()
+        )
+
+    # Note: We use reduction=sum while computing the loss.
+    info["loss"] = loss.detach().cpu().item()
+
+    return loss, info
+
+
+def compute_validation_loss(
+    params: AttributeDict,
+    model: Union[nn.Module, DDP],
+    valid_dl: torch.utils.data.DataLoader,
+    world_size: int = 1,
+) -> MetricsTracker:
+    """Run the validation process."""
+    model.eval()
+
+    tot_loss = MetricsTracker()
+
+    for batch_idx, batch in enumerate(valid_dl):
+        loss, loss_info = compute_loss(
+            params=params,
+            model=model,
+            batch=batch,
+            is_training=False,
+        )
+        assert loss.requires_grad is False
+        tot_loss = tot_loss + loss_info
+
+    if world_size > 1:
+        tot_loss.reduce(loss.device)
+
+    loss_value = tot_loss["loss"] / tot_loss["frames"]
+    if loss_value < params.best_valid_loss:
+        params.best_valid_epoch = params.cur_epoch
+        params.best_valid_loss = loss_value
+
+    return tot_loss
+
+
+def train_one_epoch(
+    params: AttributeDict,
+    model: Union[nn.Module, DDP],
+    optimizer: torch.optim.Optimizer,
+    scheduler: LRSchedulerType,
+    train_dl: torch.utils.data.DataLoader,
+    valid_dl: torch.utils.data.DataLoader,
+    scaler: GradScaler,
+    model_avg: Optional[nn.Module] = None,
+    tb_writer: Optional[SummaryWriter] = None,
+    world_size: int = 1,
+    rank: int = 0,
+) -> None:
+    """Train the model for one epoch.
+
+    The training loss from the mean of all frames is saved in
+    `params.train_loss`. It runs the validation process every
+    `params.valid_interval` batches.
+
+    Args:
+      params:
+        It is returned by :func:`get_params`.
+      model:
+        The model for training.
+      optimizer:
+        The optimizer we are using.
+      scheduler:
+        The learning rate scheduler, we call step() every step.
+      train_dl:
+        Dataloader for the training dataset.
+      valid_dl:
+        Dataloader for the validation dataset.
+      scaler:
+        The scaler used for mix precision training.
+      model_avg:
+        The stored model averaged from the start of training.
+      tb_writer:
+        Writer to write log messages to tensorboard.
+      world_size:
+        Number of nodes in DDP training. If it is 1, DDP is disabled.
+      rank:
+        The rank of the node in DDP training. If no DDP is used, it should
+        be set to 0.
+    """
+    model.train()
+
+    tot_loss = MetricsTracker()
+
+    cur_batch_idx = params.get("cur_batch_idx", 0)
+
+    saved_bad_model = False
+    def save_bad_model(suffix: str = ""):
+        save_checkpoint_impl(filename=params.exp_dir / f"bad-model{suffix}-{rank}.pt",
+                             model=model,
+                             model_avg=model_avg,
+                             params=params,
+                             optimizer=optimizer,
+                             scheduler=scheduler,
+                             scaler=scaler,
+                             rank=0)
+
+
+    for batch_idx, batch in enumerate(train_dl):
+        if batch_idx % 10 == 0:
+            set_batch_count(model, get_adjusted_batch_count(params))
+        if batch_idx < cur_batch_idx:
+            continue
+        cur_batch_idx = batch_idx
+
+        params.batch_idx_train += 1
+
+        try:
+            with torch.cuda.amp.autocast(enabled=params.use_fp16):
+                loss, loss_info = compute_loss(
+                    params=params,
+                    model=model,
+                    batch=batch,
+                    is_training=True,
+                )
+            # summary stats
+            tot_loss = (tot_loss * (1 - 1 / params.reset_interval)) + loss_info
+
+            # NOTE: We use reduction==sum and loss is computed over utterances
+            # in the batch and there is no normalization to it so far.
+            scaler.scale(loss).backward()
+            scheduler.step_batch(params.batch_idx_train)
+
+            scaler.step(optimizer)
+            scaler.update()
+            optimizer.zero_grad()
+        except:  # noqa
+            save_bad_model()
+            display_and_save_batch(batch, params=params)
+            raise
+
+        if params.print_diagnostics and batch_idx == 5:
+            return
+
+        if (
+            rank == 0
+            and params.batch_idx_train > 0
+            and params.batch_idx_train % params.average_period == 0
+        ):
+            update_averaged_model(
+                params=params,
+                model_cur=model,
+                model_avg=model_avg,
+            )
+
+        if (
+            params.batch_idx_train > 0
+            and params.batch_idx_train % params.save_every_n == 0
+        ):
+            params.cur_batch_idx = batch_idx
+            save_checkpoint_with_global_batch_idx(
+                out_dir=params.exp_dir,
+                global_batch_idx=params.batch_idx_train,
+                model=model,
+                model_avg=model_avg,
+                params=params,
+                optimizer=optimizer,
+                scheduler=scheduler,
+                scaler=scaler,
+                rank=rank,
+            )
+            del params.cur_batch_idx
+            remove_checkpoints(
+                out_dir=params.exp_dir,
+                topk=params.keep_last_k,
+                rank=rank,
+            )
+
+        if batch_idx % 100 == 0 and params.use_fp16:
+            # If the grad scale was less than 1, try increasing it.    The _growth_interval
+            # of the grad scaler is configurable, but we can't configure it to have different
+            # behavior depending on the current grad scale.
+            cur_grad_scale = scaler._scale.item()
+
+            if cur_grad_scale < 8.0 or (cur_grad_scale < 32.0 and batch_idx % 400 == 0):
+                scaler.update(cur_grad_scale * 2.0)
+            if cur_grad_scale < 0.01:
+                if not saved_bad_model:
+                    save_bad_model(suffix="-first-warning")
+                    saved_bad_model = True
+                logging.warning(f"Grad scale is small: {cur_grad_scale}")
+            if cur_grad_scale < 1.0e-05:
+                save_bad_model()
+                raise RuntimeError(f"grad_scale is too small, exiting: {cur_grad_scale}")
+
+        if batch_idx % params.log_interval == 0:
+            cur_lr = max(scheduler.get_last_lr())
+            cur_grad_scale = scaler._scale.item() if params.use_fp16 else 1.0
+
+            logging.info(
+                f"Epoch {params.cur_epoch}, "
+                f"batch {batch_idx}, loss[{loss_info}], "
+                f"tot_loss[{tot_loss}], "
+                f"lr: {cur_lr:.2e}, " +
+                (f"grad_scale: {scaler._scale.item()}" if params.use_fp16 else "")
+            )
+
+            if tb_writer is not None:
+                tb_writer.add_scalar(
+                    "train/learning_rate", cur_lr, params.batch_idx_train
+                )
+
+                loss_info.write_summary(
+                    tb_writer, "train/current_", params.batch_idx_train
+                )
+                tot_loss.write_summary(
+                    tb_writer, "train/tot_", params.batch_idx_train
+                )
+                if params.use_fp16:
+                    tb_writer.add_scalar(
+                        "train/grad_scale", cur_grad_scale, params.batch_idx_train
+                    )
+
+
+
+        if batch_idx % params.valid_interval == 0 and not params.print_diagnostics:
+            logging.info("Computing validation loss")
+            valid_info = compute_validation_loss(
+                params=params,
+                model=model,
+                valid_dl=valid_dl,
+                world_size=world_size,
+            )
+            model.train()
+            logging.info(f"Epoch {params.cur_epoch}, validation: {valid_info}")
+            logging.info(f"Maximum memory allocated so far is {torch.cuda.max_memory_allocated()//1000000}MB")
+            if tb_writer is not None:
+                valid_info.write_summary(
+                    tb_writer, "train/valid_", params.batch_idx_train
+                )
+
+    loss_value = tot_loss["loss"] / tot_loss["frames"]
+    params.train_loss = loss_value
+    if params.train_loss < params.best_train_loss:
+        params.best_train_epoch = params.cur_epoch
+        params.best_train_loss = params.train_loss
+
+
+def run(rank, world_size, args):
+    """
+    Args:
+      rank:
+        It is a value between 0 and `world_size-1`, which is
+        passed automatically by `mp.spawn()` in :func:`main`.
+        The node with rank 0 is responsible for saving checkpoint.
+      world_size:
+        Number of GPUs for DDP training.
+      args:
+        The return value of get_parser().parse_args()
+    """
+    params = get_params()
+    params.update(vars(args))
+
+    fix_random_seed(params.seed)
+    if world_size > 1:
+        setup_dist(rank, world_size, params.master_port)
+
+    setup_logger(f"{params.exp_dir}/log/log-train")
+    logging.info("Training started")
+
+    if args.tensorboard and rank == 0:
+        tb_writer = SummaryWriter(log_dir=f"{params.exp_dir}/tensorboard")
+    else:
+        tb_writer = None
+
+    device = torch.device("cpu")
+    if torch.cuda.is_available():
+        device = torch.device("cuda", rank)
+    logging.info(f"Device: {device}")
+
+    logging.info(params)
+
+    logging.info("About to create model")
+    model = get_model(params)
+
+    num_param = sum([p.numel() for p in model.parameters()])
+    logging.info(f"Number of model parameters: {num_param}")
+
+    assert params.save_every_n >= params.average_period
+    model_avg: Optional[nn.Module] = None
+    if rank == 0:
+        # model_avg is only used with rank 0
+        model_avg = copy.deepcopy(model).to(torch.float64)
+
+    assert params.start_epoch > 0, params.start_epoch
+    checkpoints = load_checkpoint_if_available(
+        params=params, model=model, model_avg=model_avg
+    )
+
+    model.to(device)
+    if world_size > 1:
+        logging.info("Using DDP")
+        model = DDP(model, device_ids=[rank],
+                    find_unused_parameters=True)
+
+    optimizer = ScaledAdam(
+        get_parameter_groups_with_lrs(
+            model, lr=params.base_lr, include_names=True),
+        lr=params.base_lr,   # should have no effect
+        clipping_scale=2.0,
+    )
+
+    scheduler = Eden(optimizer, params.lr_batches, params.lr_epochs)
+
+    if checkpoints and "optimizer" in checkpoints:
+        logging.info("Loading optimizer state dict")
+        optimizer.load_state_dict(checkpoints["optimizer"])
+
+    if (
+        checkpoints
+        and "scheduler" in checkpoints
+        and checkpoints["scheduler"] is not None
+    ):
+        logging.info("Loading scheduler state dict")
+        scheduler.load_state_dict(checkpoints["scheduler"])
+
+    if params.print_diagnostics:
+        opts = diagnostics.TensorDiagnosticOptions(
+            2 ** 22
+        )  # allow 4 megabytes per sub-module
+        diagnostic = diagnostics.attach_diagnostics(model, opts)
+
+    if params.inf_check:
+        register_inf_check_hooks(model)
+
+
+    train = LmDataset(params.train_file_list,
+                      bytes_per_segment=params.bytes_per_segment)
+    train_dl = LmDataloader(train, batch_size=params.batch_size,
+                            num_workers=params.num_workers)
+
+    valid = LmDataset(params.valid_file_list,
+                      bytes_per_segment=params.bytes_per_segment)
+    valid_dl = LmDataloader(valid, batch_size=params.batch_size,
+                            num_workers=params.num_workers)
+
+
+    scaler = GradScaler(enabled=params.use_fp16,
+                        init_scale=1.0)
+    if checkpoints and "grad_scaler" in checkpoints:
+        logging.info("Loading grad scaler state dict")
+        scaler.load_state_dict(checkpoints["grad_scaler"])
+
+    for epoch in range(params.start_epoch, params.num_epochs + 1):
+        scheduler.step_epoch(epoch - 1)
+        fix_random_seed(params.seed + epoch - 1)
+        # the above will affect random seeds in the dataloaders.
+
+        if tb_writer is not None:
+            tb_writer.add_scalar("train/epoch", epoch, params.batch_idx_train)
+
+        params.cur_epoch = epoch
+
+        train_one_epoch(
+            params=params,
+            model=model,
+            model_avg=model_avg,
+            optimizer=optimizer,
+            scheduler=scheduler,
+            train_dl=train_dl,
+            valid_dl=valid_dl,
+            scaler=scaler,
+            tb_writer=tb_writer,
+            world_size=world_size,
+            rank=rank,
+        )
+
+        if params.print_diagnostics:
+            diagnostic.print_diagnostics()
+            break
+
+        save_checkpoint(
+            params=params,
+            model=model,
+            model_avg=model_avg,
+            optimizer=optimizer,
+            scheduler=scheduler,
+            scaler=scaler,
+            rank=rank,
+        )
+
+    logging.info("Done!")
+
+    if world_size > 1:
+        torch.distributed.barrier()
+        cleanup_dist()
+
+
+def display_and_save_batch(
+    batch: Tensor,
+    params: AttributeDict,
+) -> None:
+    """Display the batch statistics and save the batch into disk.
+
+    Args:
+      batch:
+        A batch of data. See `lhotse.dataset.K2SpeechRecognitionDataset()`
+        for the content in it.
+      params:
+        Parameters for training. See :func:`get_params`.
+    """
+    from lhotse.utils import uuid4
+
+    filename = f"{params.exp_dir}/batch-{uuid4()}.pt"
+    logging.info(f"Saving batch to {filename}")
+    torch.save({'labels': batch}, filename)
+
+
+
+
+def main():
+    parser = get_parser()
+    args = parser.parse_args()
+
+    args.exp_dir = Path(args.exp_dir)
+
+    world_size = args.world_size
+    assert world_size >= 1
+    if world_size > 1:
+        mp.spawn(run, args=(world_size, args), nprocs=world_size, join=True)
+    else:
+        run(rank=0, world_size=1, args=args)
+
+
+torch.set_num_threads(1)
+torch.set_num_interop_threads(1)
+
+if __name__ == "__main__":
+    main()
diff --git a/egs/libriheavy/LM/zipformer1/zipformer.py b/egs/libriheavy/LM/zipformer1/zipformer.py
new file mode 120000
index 000000000..d053ea6de
--- /dev/null
+++ b/egs/libriheavy/LM/zipformer1/zipformer.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/zipformer2/zipformer.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py b/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py
index 8538a3cfe..bff07d8ea 100644
--- a/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py
@@ -111,8 +111,8 @@ class Zipformer2(EncoderInterface):
             dropout: FloatLike = None,  # see code below for default
             warmup_batches: float = 4000.0,
             causal: bool = False,
-            chunk_size: Tuple[int] = [-1],
-            left_context_frames: Tuple[int] = [-1],
+            chunk_size: Tuple[int] = (-1,),
+            left_context_frames: Tuple[int] = (-1,),
     ) -> None:
         super(Zipformer2, self).__init__()
 
diff --git a/egs/librispeech/ASR/zipformer2 b/egs/librispeech/ASR/zipformer2
new file mode 120000
index 000000000..fd3d16498
--- /dev/null
+++ b/egs/librispeech/ASR/zipformer2
@@ -0,0 +1 @@
+pruned_transducer_stateless7
\ No newline at end of file