AudioMNIST: forward() got an unexpected keyword argument 'lens' #2591

(py39_speechbrain) zh@zh-ThinkStation-P360-Tower:~/speechbrain-SE$ python /home/zh/speechbrain-SE/recipes/AudioMNIST/diffusion/train.py /home/zh/speechbrain-SE/recipes/AudioMNIST/diffusion/hparams/train.yaml --data_folder=/home/zh/speechbrain-SE/recipes/AudioMNIST
Could not enable Tensorboard logging - Tensorboard is not available
speechbrain.core - Beginning experiment!
speechbrain.core - Experiment folder: ./results/diffusion/baseline/1986
audiomnist_prepare - Data_preparation...
Downloading https://www.dropbox.com/scl/fi/ekibujzmvakufvm31ptrf/audiominist-meta.zip?rlkey=69vwmqcoc1xl7t5j94yjilxoc&dl=1 to /home/zh/speechbrain-SE/recipes/AudioMNIST/metadata/metadata.zip
audiominist-meta.zip?rlkey=69vwmqcoc1xl7t5j94yjilxoc&dl=1: 73.7kB [00:02, 25.0kB/s]
Extracting /home/zh/speechbrain-SE/recipes/AudioMNIST/metadata/metadata.zip to /home/zh/speechbrain-SE/recipes/AudioMNIST/metadata
audiomnist_prepare - Converting split train
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 28519/28519 [00:22<00:00, 1282.56it/s]
audiomnist_prepare - Saving metadata to ./results/diffusion/baseline/1986/save/train.json
audiomnist_prepare - Converting split valid
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 750/750 [00:01<00:00, 650.65it/s]
audiomnist_prepare - Saving metadata to ./results/diffusion/baseline/1986/save/valid.json
audiomnist_prepare - Converting split test
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 750/750 [00:01<00:00, 670.96it/s]
audiomnist_prepare - Saving metadata to ./results/diffusion/baseline/1986/save/test.json
speechbrain.core - Info: max_grad_norm arg from hparam file is used
speechbrain.core - Info: ckpt_interval_minutes arg from hparam file is used
speechbrain.core - Gradscaler enabled: False. Using precision: fp32.
speechbrain.core - DiffusionBrain Model Statistics:

• Total Number of Trainable Parameters: 398.9M
• Total Number of Parameters: 398.9M
• Trainable Parameters represent 100.0000% of the total size.
  speechbrain.utils.checkpoints - Would load a checkpoint here, but none found yet.
  speechbrain.utils.epoch_loop - Going into epoch 1
  speechbrain.utils.fetching - Fetch hyperparams.yaml: Using existing file/symlink in pretrained_models/HIFIGAN-92ca03b8a61b04650be053125e9597a9/hyperparams.yaml.
  speechbrain.utils.fetching - Fetch custom.py: Delegating to Huggingface hub, source speechbrain/tts-hifigan-libritts-16kHz.
  speechbrain.utils.fetching - Fetch generator.ckpt: Using existing file/symlink in pretrained_models/HIFIGAN-92ca03b8a61b04650be053125e9597a9/generator.ckpt.
  speechbrain.utils.parameter_transfer - Loading pretrained files for: generator
  0%| | 0/1783 [00:00<?, ?it/s]
  speechbrain.core - Exception:
  Traceback (most recent call last):
  File "/home/zh/speechbrain-SE/recipes/AudioMNIST/diffusion/train.py", line 1648, in
  diffusion_brain.fit(
  File "/home/zh/speechbrain-develop0519/speechbrain/core.py", line 1607, in fit
  self._fit_train(train_set=train_set, epoch=epoch, enable=enable)
  File "/home/zh/speechbrain-develop0519/speechbrain/core.py", line 1432, in _fit_train
  loss = self.fit_batch(batch)
  File "/home/zh/speechbrain-SE/recipes/AudioMNIST/diffusion/train.py", line 307, in fit_batch
  outputs = self.compute_forward(batch, sb.Stage.TRAIN)
  File "/home/zh/speechbrain-SE/recipes/AudioMNIST/diffusion/train.py", line 180, in compute_forward
  pred, noise, noisy_sample = self.modules.diffusion.train_sample(
  File "/home/zh/speechbrain-develop0519/speechbrain/nnet/diffusion.py", line 102, in train_sample
  pred = self.model(noisy_sample, timesteps, **kwargs)
  File "/home/zh/anaconda3/envs/py39_speechbrain/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
  return forward_call(*args, **kwargs)
  TypeError: forward() got an unexpected keyword argument 'lens'
  (py39_speechbrain) zh@zh-ThinkStation-P360-Tower:~/speechbrain-SE$

Hi @flexthink
Seeing that your answers are very professional, so @you. Could you please help me with my question? Thank you