diff --git a/research/pate_2017/input.py b/research/pate_2017/input.py
index c2d669e..bcdb6d0 100644
--- a/research/pate_2017/input.py
+++ b/research/pate_2017/input.py
@@ -17,8 +17,10 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import cPickle
-#import pickle #for python 3
+try:
+  import cPickle as pickle
+except:
+  import pickle
 import gzip
 import math
 import numpy as np
@@ -156,8 +158,7 @@ def unpickle_cifar_dic(file):
   :return: tuple of (images, labels)
   """
   fo = open(file, 'rb')
-  dict = cPickle.load(fo)
-  #dict = pickle.load(fo) #for python 3
+  dict = pickle.load(fo)
   fo.close()
   return dict['data'], dict['labels']