diff --git a/main.py b/main.py
index eadbe1f..a6fadc1 100644
--- a/main.py
+++ b/main.py
@@ -17,6 +17,7 @@
 flags.DEFINE_string("dataset", "celebA", "The name of dataset [celebA, mnist, lsun]")
 flags.DEFINE_string("checkpoint_dir", "checkpoint", "Directory name to save the checkpoints [checkpoint]")
 flags.DEFINE_string("sample_dir", "samples", "Directory name to save the image samples [samples]")
+flags.DEFINE_string("gpu", "0", "The ID of GPU for training [0]")
 flags.DEFINE_boolean("is_train", False, "True for training, False for testing [False]")
 flags.DEFINE_boolean("is_crop", False, "True for training, False for testing [False]")
 flags.DEFINE_boolean("visualize", False, "True for visualizing, False for nothing [False]")
@@ -30,7 +31,7 @@ def main(_):
     if not os.path.exists(FLAGS.sample_dir):
         os.makedirs(FLAGS.sample_dir)
 
-    with tf.Session() as sess:
+    with tf.Session(config=tf.ConfigProto(gpu_options = tf.GPUOptions(visible_device_list=FLAGS.gpu))) as sess:
         if FLAGS.dataset == 'mnist':
             dcgan = DCGAN(sess, image_size=FLAGS.image_size, batch_size=FLAGS.batch_size, y_dim=10,
                     dataset_name=FLAGS.dataset, is_crop=FLAGS.is_crop, checkpoint_dir=FLAGS.checkpoint_dir)
diff --git a/model.py b/model.py
index 0153090..7bce0aa 100755
--- a/model.py
+++ b/model.py
@@ -164,8 +164,9 @@ def generator(self, z):
         h1 = lrelu(self.h1)
 
         h2, self.h2_w, self.h2_b = deconv2d(h1, [self.batch_size, 32, 32, 3*16], d_h=1, d_w=1, name='g_h2', with_w=True)
-        h2 = PS(h2, 4, color=True)
-
+        #h2 = PS(h2, 4, color=True)
+        h2 = tf.depth_to_space(h2, 4)
+        
         return tf.nn.tanh(h2)
 
     def save(self, checkpoint_dir, step):