Why multi-gpu is slower? #2842

asr_brain.fit(
    asr_brain.hparams.epoch_counter,
    train_data,
    valid_data,
    train_loader_kwargs=hparams["train_loader_kwargs"],
    valid_loader_kwargs=hparams["valid_loader_kwargs"],
)

asr_brain.fit(
    asr_brain.hparams.epoch_counter,
    train_dataloader,
    valid_dataloader
)

#!/usr/bin/env python3
"""Recipe for training a whisper-based ASR system with CommonVoice.
The system employs whisper from OpenAI (https://cdn.openai.com/papers/whisper.pdf).
This recipe take the whisper encoder-decoder to fine-tune on.

To run this recipe, do the following:
> python train_with_whisper.py hparams/train_hf_whisper.yaml

Authors
 * Pooneh Mousavi 2022
 * Adel Moumen 2024
"""

import logging
import os
import sys

import torch
from hyperpyyaml import load_hyperpyyaml

import speechbrain as sb
from speechbrain.utils.distributed import run_on_main
from vbox_custom.dataset import AudioTextDataset
from vbox_custom.core import ASR
from vbox_custom.dataloader import make_multi_dataloader


logger = logging.getLogger(__name__)


os.environ["TOKENIZERS_PARALLELISM"] = "false"


def dataio_prepare(hparams, tokenizer, device="cuda"):
    """This function prepares the datasets (or data loaders) to be used in the brain class."""
    datasets_config = hparams["datasets"]
    train_csv_paths = {ds_name: f'{hparams["save_folder"]}/train_{ds_name}.csv' for ds_name in datasets_config}
    valid_csv_paths = {ds_name: f'{hparams["save_folder"]}/valid_{ds_name}.csv' for ds_name in datasets_config}
    test_csv_paths = {ds_name: f'{hparams["save_folder"]}/test_{ds_name}.csv' for ds_name in datasets_config}
    data_folders = {ds_name: config['data_folder'] for ds_name, config in datasets_config.items()}
    cache_dir = f'{hparams["save_folder"]}/cache_vec'

    # Create datasets
    train_data = {
        ds_name: AudioTextDataset(
            csv_path=train_csv_paths[ds_name],
            data_folder=data_folders[ds_name],
            tokenizer=tokenizer,
            hparams=hparams,
            device=device,
            sorting=hparams["sorting"],
            avoid_if_longer_than=hparams["avoid_if_longer_than"],
            cache_dir=cache_dir,
        ) for ds_name in datasets_config
    }

    valid_data = {
        ds_name: AudioTextDataset(
            csv_path=valid_csv_paths[ds_name],
            data_folder=data_folders[ds_name],
            tokenizer=tokenizer,
            hparams=hparams,
            device=device,
            sorting="ascending",  # Always sort validation data ascending
            avoid_if_longer_than=hparams["avoid_if_longer_than"],
            cache_dir=cache_dir,
        ) for ds_name in datasets_config
    }

    test_data = {
        ds_name: AudioTextDataset(
            csv_path=test_csv_paths[ds_name],
            data_folder=data_folders[ds_name],
            tokenizer=tokenizer,
            hparams=hparams,
            device=device,
            sorting="ascending",  # Always sort test data ascending
            avoid_if_longer_than=hparams["avoid_if_longer_than"],
            cache_dir=cache_dir,
        ) for ds_name in datasets_config
    }
    
    return train_data, valid_data, test_data


def collect_linear(model):
    """Returns a list of linear layers from a given model
    Arguments
    ---------
    model : torch.nn.Module
        A torch model to collect the layer names from.
    Returns
    -------
    layers : list
        A list of layer names for linear layers.
    """
    layers = []
    for layer_name, layer in model.named_modules():
        if isinstance(layer, torch.nn.Linear):
            layers.append(layer_name)
    return layers


if __name__ == "__main__":
    # CLI:
    hparams_file, run_opts, overrides = sb.parse_arguments(sys.argv[1:])

    # create ddp_group with the right communication protocol
    sb.utils.distributed.ddp_init_group(run_opts)

    with open(hparams_file) as fin:
        hparams = load_hyperpyyaml(fin, overrides)

    # Create experiment directory
    sb.create_experiment_directory(
        experiment_directory=hparams["output_folder"],
        hyperparams_to_save=hparams_file,
        overrides=overrides,
    )

    # Setting the save folder
    os.makedirs(hparams["save_folder"], exist_ok=True)

    # Dataset prep
    from recipes.VBoxMultiple.prepare import prepare_ds

    datasets_config = hparams["datasets"]
    train_csv_paths = {ds_name: f'{hparams["save_folder"]}/train_{ds_name}.csv' for ds_name in datasets_config}
    valid_csv_paths = {ds_name: f'{hparams["save_folder"]}/valid_{ds_name}.csv' for ds_name in datasets_config}
    test_csv_paths = {ds_name: f'{hparams["save_folder"]}/test_{ds_name}.csv' for ds_name in datasets_config}
    data_folders = {ds_name: config['data_folder'] for ds_name, config in datasets_config.items()}

    for ds_name in datasets_config:
        run_on_main(
            prepare_ds,
            kwargs={
                "data_folder": data_folders[ds_name],
                "save_csv_train": train_csv_paths[ds_name],
                "save_csv_valid": valid_csv_paths[ds_name],
                "save_csv_test": test_csv_paths[ds_name],
                "split_ratios": hparams["split_ratio"],
                "transcript_file": f"{data_folders[ds_name]}/metadata.csv",
                "accented_letters": hparams["accented_letters"],
                "language": hparams["language"],
                "skip_prep": hparams["skip_prep"],
            },
        )

    # Defining tokenizer and loading it
    tokenizer = hparams["whisper"].tokenizer

    # here we create the datasets objects as well as tokenization and encoding
    train_data, valid_data, test_data = dataio_prepare(hparams, tokenizer)
    train_data_list = []
    valid_data_list = []
    test_data_list = []
    train_ratios = []
    valid_ratios = []
    for ds_name in datasets_config:
        train_data_list.append(train_data[ds_name])
        valid_data_list.append(valid_data[ds_name])
        test_data_list.append(test_data[ds_name])
        train_ratios.append(hparams['dataset_ratios']['train'][ds_name])
        valid_ratios.append(hparams['dataset_ratios']['valid'][ds_name])

    train_dataloader = make_multi_dataloader(train_data_list,
                                        ratios=train_ratios,
                                        samples=hparams["sample_size"]["train"],
                                        batch_size=hparams["train_loader_kwargs"]["batch_size"],
                                        device=run_opts["device"])

    valid_dataloader = make_multi_dataloader(valid_data_list,
                                        ratios=valid_ratios,
                                        samples=hparams["sample_size"]["valid"],
                                        batch_size=hparams["valid_loader_kwargs"]["batch_size"],
                                        device=run_opts["device"])
    # train_data_concat = torch.utils.data.ConcatDataset(train_data_list)
    # valid_data_concat = torch.utils.data.ConcatDataset(valid_data_list)


    # We load the pretrained whisper model
    if "pretrainer" in hparams.keys():
        run_on_main(hparams["pretrainer"].collect_files)
        hparams["pretrainer"].load_collected(run_opts["device"])
    
    # Add LoRA to the pretrained model
    if "lora_config" in hparams:
        import peft
        # lora_layers = collect_linear(hparams["whisper"])
        # model_config = hparams["lora_config"](target_modules=lora_layers)
        model_config = hparams["lora_config"]()
        hparams["whisper"] = peft.get_peft_model(
            model=hparams["whisper"], peft_config=model_config
        )

    # Trainer initialization
    asr_brain = ASR(
        modules=hparams["modules"],
        hparams=hparams,
        run_opts=run_opts,
        checkpointer=hparams["checkpointer"],
        opt_class=hparams["whisper_opt_class"],
    )

    # We dynamically add the tokenizer to our brain class.
    # NB: This tokenizer corresponds to the one used for Whisper.
    asr_brain.tokenizer = tokenizer

    # Training
    asr_brain.fit(
        asr_brain.hparams.epoch_counter,
        train_dataloader,
        valid_dataloader
    )

    # Testing
    asr_brain.hparams.test_wer_file = hparams["test_wer_file"]
    for ds_name in datasets_config:
        asr_brain.evaluate(
            test_data[ds_name],
            min_key="WER",
            test_loader_kwargs=hparams["test_loader_kwargs"],
        )

import os

import torch
import torchaudio
from torch.utils.data import Dataset
from tqdm import tqdm

import speechbrain as sb

class AudioTextDataset(Dataset):
    def __init__(
        self,
        csv_path,
        data_folder,
        tokenizer,
        hparams,
        device="cuda",
        sorting="random",
        avoid_if_longer_than=None,
        cache_dir="cached_signals",  # Directory to store preprocessed signals
    ):
        self.data_folder = data_folder
        self.tokenizer = tokenizer
        self.hparams = hparams
        self.device = device
        self.sorting = sorting
        self.avoid_if_longer_than = avoid_if_longer_than
        self.cache_dir = cache_dir

        # Create cache directory if it doesn't exist
        os.makedirs(self.cache_dir, exist_ok=True)

        # Load data from CSV
        self.data = []
        with open(csv_path, "r") as f:
            lines = f.readlines()
            header = lines[0].strip().split(",")
            for line in lines[1:]:
                values = line.strip().split(",")
                sample = dict(zip(header, values))
                self.data.append(sample)

        # Preprocess all samples
        self.preprocess_data()

        # Filter and sort data
        self.filter_and_sort()

    def preprocess_data(self):
        """Preprocess all samples (audio and text) and store them."""
        for sample in tqdm(self.data, desc="Preprocessing data"):
            file_name = sample["wav"].replace("/", "_").replace("\\", "_")
            cache_path = os.path.join(self.cache_dir, f"{file_name}.pt")
            sample["sig_cache_path"] = cache_path  # Store path to cached signal

            if not os.path.exists(cache_path):
                # Preprocess audio
                wav_path = sample["wav"]
                info = torchaudio.info(wav_path)
                sig = sb.dataio.dataio.read_audio(wav_path)
                if info.sample_rate != self.hparams["sample_rate"]:
                    resampler = torchaudio.transforms.Resample(
                        info.sample_rate, self.hparams["sample_rate"]
                    )
                    sig = resampler(sig)
                if sig.ndim > 1:
                    sig = torch.mean(sig, dim=0)

                # Save preprocessed signal to disk
                torch.save(sig, cache_path)
            
            # sample["duration"] = len(sig) / self.hparams["sample_rate"]  # Add duration
            sample["duration"] = float(sample["duration"])

            # Preprocess text
            wrd = sample["wrd"]
            if self.hparams.get("normalized_transcripts", False):
                wrd = self.tokenizer.normalize(wrd)
            tokens_list = self.tokenizer.encode(wrd, add_special_tokens=False)
            tokens_list = self.tokenizer.build_inputs_with_special_tokens(tokens_list)
            sample["tokens_list"] = tokens_list
            sample["tokens_bos"] = torch.LongTensor(tokens_list[:-1])
            sample["tokens_eos"] = torch.LongTensor(tokens_list[1:])
            sample["tokens"] = torch.LongTensor(tokens_list)

    def filter_and_sort(self):
        """Filter and sort the dataset based on duration."""
        # Filter samples based on duration
        if self.avoid_if_longer_than is not None:
            self.data = [
                sample
                for sample in self.data
                if sample["duration"] <= self.avoid_if_longer_than
            ]

        # Sort samples based on duration
        if self.sorting == "ascending":
            self.data.sort(key=lambda x: x["duration"])
        elif self.sorting == "descending":
            self.data.sort(key=lambda x: x["duration"], reverse=True)
        elif self.sorting == "random":
            pass
        else:
            raise ValueError("sorting must be 'random', 'ascending', or 'descending'")

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        sample = self.data[idx]
        # Load preprocessed signal from disk
        sig = torch.load(sample["sig_cache_path"])
        return {
            "id": sample["ID"],
            "sig": sig,
            "tokens_list": sample["tokens_list"],
            "tokens_bos": sample["tokens_bos"],
            "tokens_eos": sample["tokens_eos"],
            "tokens": sample["tokens"],
        }

import math
from typing import List
from torch.utils.data import Dataset, DataLoader, RandomSampler, IterableDataset
import torch
from speechbrain.dataio.batch import PaddedBatch
from speechbrain.dataio.sampler import ReproducibleRandomSampler
from speechbrain.dataio.dataloader import SaveableDataLoader, LoopedLoader
from vbox_custom.batch import VBoxPaddedBatch


class CombinedDataLoader(DataLoader):
    def __init__(self, dataloaders : List[DataLoader], num_batches, device='cuda'):
        """
        Initializes the combined dataloader.

        Args:
            dataloaders: A list of DataLoader instances.
            num_batches: Total number of batches to draw from the combined loader.
        """
        self.dataloaders = dataloaders
        self.num_batches = num_batches
        self.current_index = 0
        self.device = device
        self.dataset = None

    def _gather_data(self):
        """
        Collects all data from the loaders into a single list.
        This makes the combined loader subscriptable.
        """
        iters = [iter(loader) for loader in self.dataloaders]

        for _ in range(self.num_batches):
            batches = [next(iterator).to(self.device) for iterator in iters]
            concat_batch = VBoxPaddedBatch.merge(batches)
            yield concat_batch

        torch.cuda.empty_cache()

    def __iter__(self):
        self.dataset = self._gather_data()
        self.iter = iter(self.dataset)
        self.current_index = 0
        return self

    def __next__(self):
        if self.current_index < self.num_batches:
            # print(f'self.current_index = {self.current_index}')
            item = next(self.iter)
            self.current_index += 1
            return item
        else:
            raise StopIteration

    def __len__(self):
        return self.num_batches


def make_multi_dataloader(datasets: List[Dataset],
                    ratios: List[float],
                    samples: int,
                    device: str = 'cuda',
                    batch_size=8):
    """
    Creates a data loader which samples from multiple datasets randomly with ratios given.
    """
    # Normalize ratios
    ratios_sum = sum(ratios)
    if ratios_sum != 1:
        ratios = [ratio / ratios_sum for ratio in ratios]
    
    # Calculate batch sizes
    batch_sizes = [max(int(batch_size * ratio), 1) for ratio in ratios[:-1]]
    batch_sizes.append(batch_size - sum(batch_sizes))
    print(f'numbers in each batch: {batch_sizes}')

    # Create DataLoaders
    dataloaders = []
    counter = 1
    for dataset, ds_batch_size in zip(datasets, batch_sizes):
        ds_samples = int(math.ceil(samples / batch_size)) * ds_batch_size
        print(f'dataset {counter}: sampling {ds_samples} from {len(dataset)}')
        
        # If ds_samples <= len(dataset), standard random sampling without replacement
        if ds_samples <= len(dataset):
            sampler = RandomSampler(dataset, replacement=False, num_samples=ds_samples)
            print("replacement=False")
        else:
            # Ensure that all samples are visited at least once
            num_full_passes = ds_samples // len(dataset)  # Number of full passes through the dataset
            remainder = ds_samples % len(dataset)  # Remainder that needs extra sampling
            
            # Create a custom sampler to ensure all samples are visited at least once
            # Step 1: Create a random permutation of the dataset
            indices = torch.randperm(len(dataset)).tolist()
            
            # Step 2: Repeat the dataset as necessary to meet ds_samples, but ensure all samples are used at least once
            sampled_indices = indices * num_full_passes + indices[:remainder]
            
            # Step 3: Create a sampler using the sampled indices
            sampler = torch.utils.data.sampler.SubsetRandomSampler(sampled_indices)
            print(f"repeating {num_full_passes} times")
            print("replacement=True")
        
        # Create DataLoader with the updated sampler
        dataloader = DataLoader(
            dataset, 
            batch_size=ds_batch_size, 
            sampler=sampler,
            collate_fn=PaddedBatch, 
            # shuffle=True, 
            num_workers=5,
            pin_memory=True,
            # pin_memory_device=device,
            )
        dataloaders.append(dataloader)
        counter += 1

    # Create the combined dataloader
    combined_loader = CombinedDataLoader(
        dataloaders,
        num_batches=int(samples / batch_size),
        device=device,
    )

    return combined_loader

import logging

from torch.utils.data import DataLoader
import torch

import speechbrain as sb
from speechbrain.dataio.dataloader import LoopedLoader, SaveableDataLoader
from speechbrain.utils.data_utils import undo_padding
from speechbrain.utils.distributed import if_main_process

import vbox_custom
import vbox_custom.dataloader

logger = logging.getLogger(__name__)


# Define training procedure
class ASR(sb.Brain):
    def compute_forward(self, batch, stage):
        """Forward computations from the waveform batches to the output probabilities."""
        batch = batch.to(self.device)
        wavs, wav_lens = batch.sig
        bos_tokens, bos_tokens_lens = batch.tokens_bos

        # Add waveform augmentation if specified.
        if stage == sb.Stage.TRAIN and hasattr(self.hparams, "wav_augment"):
            wavs, wav_lens = self.hparams.wav_augment(wavs, wav_lens)
            bos_tokens = self.hparams.wav_augment.replicate_labels(bos_tokens)
            bos_tokens_lens = self.hparams.wav_augment.replicate_labels(
                bos_tokens_lens
            )

        # We compute the padding mask and replace the values with the pad_token_id
        # that the Whisper decoder expect to see.
        abs_tokens_lens = bos_tokens_lens * bos_tokens.shape[1]
        pad_mask = (
                torch.arange(abs_tokens_lens.max(), device=self.device)[None, :]
                < abs_tokens_lens[:, None]
        )
        bos_tokens[~pad_mask] = self.tokenizer.pad_token_id

        # Forward encoder + decoder
        # logger.info(str(self.modules.whisper))
        enc_out, logits, _ = self.modules.whisper(wavs, bos_tokens)
        log_probs = self.hparams.log_softmax(logits)

        hyps = None
        if stage == sb.Stage.VALID:
            hyps, _, _, _ = self.hparams.valid_search(
                enc_out.detach(), wav_lens
            )
        elif stage == sb.Stage.TEST:
            hyps, _, _, _ = self.hparams.test_search(enc_out.detach(), wav_lens)

        return log_probs, hyps, wav_lens

    def compute_objectives(self, predictions, batch, stage):
        """Computes the loss NLL given predictions and targets."""

        (log_probs, hyps, wav_lens) = predictions
        batch = batch.to(self.device)
        ids = batch.id
        tokens_eos, tokens_eos_lens = batch.tokens_eos

        # Augment Labels
        if stage == sb.Stage.TRAIN and hasattr(self.hparams, "wav_augment"):
            tokens_eos = self.hparams.wav_augment.replicate_labels(tokens_eos)
            tokens_eos_lens = self.hparams.wav_augment.replicate_labels(
                tokens_eos_lens
            )

        loss = self.hparams.nll_loss(
            log_probs, tokens_eos, length=tokens_eos_lens
        )

        if stage != sb.Stage.TRAIN:
            tokens, tokens_lens = batch.tokens

            # Decode token terms to words
            predicted_words = [
                self.tokenizer.decode(t, skip_special_tokens=True).strip()
                for t in hyps
            ]

            # Convert indices to words
            target_words = undo_padding(tokens, tokens_lens)
            target_words = self.tokenizer.batch_decode(
                target_words, skip_special_tokens=True
            )

            if hasattr(self.hparams, "normalized_transcripts"):
                predicted_words = [
                    self.tokenizer.normalize(text).split(" ")
                    for text in predicted_words
                ]

                target_words = [
                    self.tokenizer.normalize(text).split(" ")
                    for text in target_words
                ]
            else:
                predicted_words = [text.split(" ") for text in predicted_words]
                target_words = [text.split(" ") for text in target_words]

            self.wer_metric.append(ids, predicted_words, target_words)
            self.cer_metric.append(ids, predicted_words, target_words)

        return loss

    def on_stage_start(self, stage, epoch):
        """Gets called at the beginning of each epoch"""
        if stage != sb.Stage.TRAIN:
            self.cer_metric = self.hparams.cer_computer()
            self.wer_metric = self.hparams.error_rate_computer()

    def on_stage_end(self, stage, stage_loss, epoch):
        """Gets called at the end of an epoch."""
        # Compute/store important stats
        stage_stats = {"loss": stage_loss}
        if stage == sb.Stage.TRAIN:
            self.train_stats = stage_stats
        else:
            stage_stats["CER"] = self.cer_metric.summarize("error_rate")
            stage_stats["WER"] = self.wer_metric.summarize("error_rate")

        # Perform end-of-iteration things, like annealing, logging, etc.
        if stage == sb.Stage.VALID:
            if hasattr(self.hparams.lr_annealing_whisper, "current_lr"):
                lr = self.hparams.lr_annealing_whisper.current_lr
                stats_meta={"epoch": epoch, "lr": lr}
            else:
                stats_meta={"epoch": epoch}
            self.hparams.train_logger.log_stats(
                stats_meta=stats_meta,
                train_stats=self.train_stats,
                valid_stats=stage_stats,
            )
            self.checkpointer.save_and_keep_only(
                meta={"WER": stage_stats["WER"]},
                min_keys=["WER"],
            )
        elif stage == sb.Stage.TEST:
            self.hparams.train_logger.log_stats(
                stats_meta={"Epoch loaded": self.hparams.epoch_counter.current},
                test_stats=stage_stats,
            )
            if if_main_process():
                with open(self.hparams.test_wer_file, "w") as w:
                    self.wer_metric.write_stats(w)

    def fit(
        self,
        epoch_counter,
        train_set,
        valid_set=None,
        progressbar=None,
        train_loader_kwargs={},
        valid_loader_kwargs={},
    ):
        if self.test_only:
            logger.info(
                "Test only mode, skipping training and validation stages."
            )
            return

        if not (
            isinstance(train_set, DataLoader)
            or isinstance(train_set, LoopedLoader)
        ):
            train_set = self.make_dataloader(
                train_set, stage=sb.Stage.TRAIN, **train_loader_kwargs
            )
        if valid_set is not None and not (
            isinstance(valid_set, DataLoader)
            or isinstance(valid_set, LoopedLoader)
        ):
            valid_set = self.make_dataloader(
                valid_set,
                stage=sb.Stage.VALID,
                ckpt_prefix=None,
                **valid_loader_kwargs,
            )

        self.on_fit_start()

        if progressbar is None:
            progressbar = not self.noprogressbar

        # Only show progressbar if requested and main_process
        enable = progressbar and sb.utils.distributed.if_main_process()

        # Iterate epochs
        for epoch in epoch_counter:
            self._fit_train(train_set=train_set, epoch=epoch, enable=enable)
            self._fit_valid(valid_set=valid_set, epoch=epoch, enable=enable)

            # Debug mode only runs a few epochs
            if (
                self.debug
                and epoch == self.debug_epochs
                or self._optimizer_step_limit_exceeded
            ):
                break

    def make_dataloader(
        self, dataset, stage, ckpt_prefix="dataloader-", **loader_kwargs
    ):
        # TRAIN stage is handled specially.
        if stage == sb.Stage.TRAIN:
            loader_kwargs = self._train_loader_specifics(dataset, loader_kwargs)
        # This commented-out code block is useful when one can ensure
        # metric reporting is DDP-valid for VALID & EVAL datasets.
        # elif self.distributed_launch:
        #     loader_kwargs = sb.dataio.dataloader.distributed_loader_specifics(
        #         self.distributed_launch, self.rank, dataset, loader_kwargs
        #     )
        dataloader = vbox_custom.dataloader.make_dataloader(
            dataset, **loader_kwargs
        )

        if (
            self.checkpointer is not None
            and ckpt_prefix is not None
            and (
                isinstance(dataloader, SaveableDataLoader)
                or isinstance(dataloader, LoopedLoader)
            )
        ):
            ckpt_key = ckpt_prefix + stage.name
            self.checkpointer.add_recoverable(ckpt_key, dataloader)
        return dataloader

import itertools
from typing import List
import torch
from speechbrain.dataio.batch import PaddedBatch, PaddedData
from speechbrain.utils.data_utils import batch_pad_right


class VBoxPaddedBatch(PaddedBatch):

    def __init__(
        self,
        examples=[],
        padded_keys=None,
        device_prep_keys=None,
        padding_func=batch_pad_right,
        padding_kwargs={},
        apply_default_convert=True,
        nonpadded_stack=True,
    ):
        self.length = len(examples)
        super().__init__(examples, padded_keys, device_prep_keys, padding_func, padding_kwargs, apply_default_convert, nonpadded_stack)

    def set_len(self, length: int):
        self.length = length

    def __len__(self):
        return self.length

    @classmethod
    def merge(cls, batches: List[PaddedBatch]):
        new_batch = cls([{}])
        list_keys = ["id", "tokens_list"]
        tensor_keys = ["sig", "tokens_bos", "tokens_eos", "tokens"]

        # Merge the attributes of type `list`
        for key in list_keys:
            setattr(new_batch, key, list(itertools.chain(*(getattr(batch, key) for batch in batches))))

        # Merge the attributes of type `tensor`. They must be padded to right.
        for key in tensor_keys:
            max_lens = [getattr(batch, key).data.shape[1] for batch in batches]
            total_max_len = max(*max_lens)
            pad_lens = [total_max_len - length for length in max_lens]
            data_list = [getattr(batch, key).data for batch in batches]
            # print(f'devices: {[data.device for data in data_list]}')
            new_data = [torch.nn.functional.pad(data, (0, pad), "constant", 0) for data, pad in zip(data_list, pad_lens)]
            new_data = torch.cat(new_data)
            lengths_list = [getattr(batch, key).lengths for batch in batches]
            new_lengths = [lengths * max_len / total_max_len for lengths, max_len in zip(lengths_list, max_lens)]
            new_lengths = torch.cat(new_lengths)
            new_value = PaddedData(data=new_data, lengths=new_lengths)
            setattr(new_batch, key, new_value)

        # Set the new batch size
        new_batch.set_len(getattr(new_batch, tensor_keys[0]).data.shape[0])
        
        return new_batch