small fix

2019-12-09 13:27:46 +01:00 · 2019-12-09 13:27:46 +01:00 · 40a533530b
parent 956f8b7672
commit 40a533530b
1 changed files with 1 additions and 1 deletions
--- a/train.py
+++ b/train.py
@ -592,7 +592,7 @@ def main(args):  # pylint: disable=redefined-outer-name
        args.restore_step = 0

    if use_cuda:
-        model = model.cuda()
+        model.cuda()
        criterion.cuda()
        if criterion_st:
            criterion_st.cuda()