update

examples/conv_tasnet/step_2_train_model.py

@@ -41,13 +41,15 @@ def get_args():
 
     parser.add_argument("--max_epochs", default=200, type=int)
 
+    parser.add_argument("--batch_size", default=64, type=int)
+    parser.add_argument("--learning_rate", default=1e-4, type=float)
     parser.add_argument("--num_serialized_models_to_keep", default=10, type=int)
     parser.add_argument("--patience", default=5, type=int)
     parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
+    parser.add_argument("--seed", default=1234, type=int)
 
     parser.add_argument("--config_file", default="config.yaml", type=str)
 
-    parser.add_argument("--seed", default=1234, type=int)
 
     args = parser.parse_args()
     return args
@@ -139,7 +141,7 @@ def main():
     )
     train_data_loader = DataLoader(
         dataset=train_dataset,
-        batch_size=config.batch_size,
+        batch_size=args.batch_size,
         # shuffle=True,
         sampler=None,
         # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
@@ -150,7 +152,7 @@ def main():
     )
     valid_data_loader = DataLoader(
         dataset=valid_dataset,
-        batch_size=config.batch_size,
+        batch_size=args.batch_size,
         # shuffle=True,
         sampler=None,
         # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.