From 0e95f6738bd1b3578a76907213dd343c8d95e731 Mon Sep 17 00:00:00 2001
From: "Sangchun Ha (Patrick)" <seomk9896@naver.com>
Date: Sun, 22 May 2022 13:40:55 +0900
Subject: [PATCH] Fix setup, test function parameter

---
 openspeech_cli/hydra_train.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/openspeech_cli/hydra_train.py b/openspeech_cli/hydra_train.py
index f717e99..9aba5ef 100644
--- 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__':