Fix bug of KG train.py script. (#922)

It cannot work when only mxnet backend is installed.

Fix bug of KG train.py script. (#922)
It cannot work when only mxnet backend is installed.
20439e1c · xiang song(charlie.song) · Da Zheng · 93e3c49d · 20439e1c
Commit 20439e1c authored Oct 12, 2019 by xiang song(charlie.song) Committed by Da Zheng Oct 11, 2019
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

apps/kg/train.py apps/kg/train.py +2 -1

No files found.
--- a/apps/kg/train.py
+++ b/apps/kg/train.py
 from dataloader import EvalDataset, TrainDataset, NewBidirectionalOneShotIterator
 from dataloader import get_dataset
-import torch.multiprocessing as mp

 import argparse
 import os
@@ -9,10 +8,12 @@ import time

 backend = os.environ.get('DGLBACKEND')
 if backend.lower() == 'mxnet':
+    import multiprocessing as mp
    from train_mxnet import load_model
    from train_mxnet import train
    from train_mxnet import test
 else:
+    import torch.multiprocessing as mp
    from train_pytorch import load_model
    from train_pytorch import train
    from train_pytorch import test