fix step_per_epoch

d9f6c64f · dengkaipeng · c3d26b92 · d9f6c64f
隐藏空白更改
内联并排

Showing with 1 addition and 1 deletion

tsm/main.py tsm/main.py +1 -1

未找到文件。
--- a/tsm/main.py
+++ b/tsm/main.py
@@ -77,7 +77,7 @@ def main():

    step_per_epoch = int(len(train_dataset) / FLAGS.batch_size \
                         / ParallelEnv().nranks)
-    optim = make_optimizer(len(train_dataset), model.parameters())
+    optim = make_optimizer(step_per_epoch, model.parameters())

    inputs = [Input([None, 8, 3, 224, 224], 'float32', name='image')]
    labels = [Input([None, 1], 'int64', name='label')]