diff --git a/face-generation/helper.py b/face-generation/helper.py
index a1a66a6..220d050 100755
--- a/face-generation/helper.py
+++ b/face-generation/helper.py
@@ -1,7 +1,10 @@
 import math
 import os
 import hashlib
-from urllib.request import urlretrieve
+try:
+    from urllib.request import urlretrieve
+except ImportError:
+    from urllib import urlretrieve
 import zipfile
 import gzip
 import shutil
@@ -211,7 +214,7 @@ def get_batches(self, batch_size):
         while current_index + batch_size <= self.shape[0]:
             data_batch = get_batch(
                 self.data_files[current_index:current_index + batch_size],
-                *self.shape[1:3],
+                self.shape[1],self.shape[2],
                 self.image_mode)
 
             current_index += batch_size
diff --git a/face-generation/problem_unittests.py b/face-generation/problem_unittests.py
index 99a3ded..cac424f 100755
--- a/face-generation/problem_unittests.py
+++ b/face-generation/problem_unittests.py
@@ -1,5 +1,8 @@
 from copy import deepcopy
-from unittest import mock
+try:
+    from unittest import mock
+except ImportError:
+    from mock import mock
 import tensorflow as tf