Fix setup, test function parameter

openspeech-team · May 22, 2022 · 0e95f67 · 0e95f67
1 parent 7e65135
commit 0e95f67
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/openspeech_cli/hydra_train.py b/openspeech_cli/hydra_train.py
@@ -22,7 +22,6 @@
 
 import os
 import hydra
-import wandb
 import pytorch_lightning as pl
 from omegaconf import DictConfig, OmegaConf
 from pytorch_lightning.utilities import rank_zero_info
@@ -45,13 +44,13 @@ def hydra_main(configs: DictConfig) -> None:
     data_module.prepare_data()
     tokenizer = TOKENIZER_REGISTRY[configs.tokenizer.unit](configs)
 
-    data_module.setup(tokenizer=tokenizer)
+    data_module.setup()
 
     model = MODEL_REGISTRY[configs.model.model_name](configs=configs, tokenizer=tokenizer)
 
     trainer = get_pl_trainer(configs, num_devices, logger)
     trainer.fit(model, data_module)
-    trainer.test()
+    trainer.test(model, data_module)
 
 
 if __name__ == '__main__':