diff --git a/gpu/train.py b/gpu/train.py
index f4c09c9..c7075ed 100755
--- a/gpu/train.py
+++ b/gpu/train.py
@@ -30,6 +30,7 @@ from torch.optim import SGD
 from torch.utils.data import DataLoader
 from torchvision.transforms import ToTensor
 import argparse
+import datetime
 
 # Training settings
 parser = argparse.ArgumentParser(description='PyTorch MNIST Example')