fix hyper param error

bs7280 · Mar 10, 2022 · 6ee791e · 6ee791e
1 parent 5902311
commit 6ee791e
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/train.py b/train.py
@@ -4,6 +4,7 @@
 
 import argparse
 import os
+from random import shuffle
 import sys
 from pathlib import Path
 
@@ -75,7 +76,11 @@ def train(
     experiment.log_parameters(hyper_params)
 
     # data generator
-    data_gen = data.DataLoader(dataset, **hyper_params)
+    data_gen = data.DataLoader(
+        dataset,
+        batch_size=hyper_params['batch_size'],
+        shuffle=hyper_params['shuffle'],
+        num_workers=hyper_params['num_workers'])
 
     # model
     base_model = model_factory[select_model]