diff --git a/src/train.py b/src/train.py
index 29f0722..413742e 100644
--- a/src/train.py
+++ b/src/train.py
@@ -25,6 +25,7 @@ import yaml
 import numpy as np
 
 import torch
+import torch.distributed as dist
 import torch.multiprocessing as mp
 import torch.nn.functional as F
 from torch.nn.parallel import DistributedDataParallel