diff --git a/research/pate_2017/input.py b/research/pate_2017/input.py
index bcdb6d0..382787e 100644
--- a/research/pate_2017/input.py
+++ b/research/pate_2017/input.py
@@ -17,10 +17,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-try:
-  import cPickle as pickle
-except:
-  import pickle
+from six.moves import pickle
 import gzip
 import math
 import numpy as np