Create DataParallel model if several GPUs

2018-11-03 10:10:01 -04:00
parent 5889765a7c
commit 5f432480c0
3 changed files with 9 additions and 0 deletions
--- a/run_classifier_pytorch.py
+++ b/run_classifier_pytorch.py
@@ -482,6 +482,9 @@ def main():
    if args.init_checkpoint is not None:
        model.bert.load_state_dict(torch.load(args.init_checkpoint, map_location='cpu'))
    model.to(device)
+    
+    if n_gpu > 1:
+        model = torch.nn.DataParallel(model)

    optimizer = BERTAdam([{'params': [p for n, p in model.named_parameters() if n != 'bias'], 'l2': 0.01},
                          {'params': [p for n, p in model.named_parameters() if n == 'bias'], 'l2': 0.}