Update train.py

This commit is contained in:
Rezakh20 2024-03-05 12:21:30 +03:30 committed by GitHub
parent 8836e004c8
commit bf5787c438
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -382,6 +382,9 @@ def get_params() -> AttributeDict:
- warm_step: The warm_step for Noam optimizer.
"""
parser = get_parser()
LibriSpeechAsrDataModule.add_arguments(parser)
args = parser.parse_args()
params = AttributeDict(
{
"best_train_loss": float("inf"),
@ -394,6 +397,7 @@ def get_params() -> AttributeDict:
"valid_interval": 800, # For the 100h subset, use 800
"alignment_interval": 25,
# parameters for conformer
"feature_dim": 768,
"subsampling_factor": 2,
"encoder_dim": 512,
"nhead": 8,