train.py: renaming function, removing dataparallel

2018-11-27 18:04:12 -08:00 · 2018-11-27 18:04:12 -08:00 · f06063f746
parent 3045ba125b
commit f06063f746
1 changed files with 1 additions and 3 deletions
--- a/train.py
+++ b/train.py
@ -84,9 +84,7 @@ def load_model(hparams):
        model.decoder.attention_layer.score_mask_value = float(finfo('float16').min)

    if hparams.distributed_run:
-        model = DistributedDataParallel(model)
-    elif torch.cuda.device_count() > 1:
-        model = DataParallel(model)
+        model = apply_gradient_allreduce(model)

    return model