diff --git a/models.py b/models.py
index 4ff615e..9fe331d 100644
--- a/models.py
+++ b/models.py
@@ -1,9 +1,9 @@
 import tensorflow as tf
 import numpy as np
 
-def get_weights(shape, mask=None):
+def get_weights(shape, name, mask=None):
     weights_initializer = tf.contrib.layers.xavier_initializer()
-    W = tf.get_variable("weights", shape, tf.float32, weights_initializer)
+    W = tf.get_variable(name, shape, tf.float32, weights_initializer)
 
     if mask:
         filter_mid_x = shape[0]//2
@@ -18,9 +18,8 @@ def get_weights(shape, mask=None):
         W *= mask_filter 
     return W
 
-
-def get_bias(shape):
-    return tf.get_variable("biases", shape, tf.float32, tf.zeros_initializer)
+def get_bias(shape, name):
+    return tf.get_variable(name, shape, tf.float32, tf.zeros_initializer)
 
 def conv_op(x, W):
     return tf.nn.conv2d(x, W, strides=[1,1,1,1], padding='SAME')
@@ -40,11 +39,11 @@ def __init__(self, W_shape, b_shape, fan_in, gated=True, payload=None, mask=None
             self.simple_conv()
 
     def gated_conv(self):
-        W_f = get_weights(self.W_shape, mask=self.mask)
-        b_f = get_bias(self.b_shape)
-        W_g = get_weights(self.W_shape, mask=self.mask)
-        b_g = get_bias(self.b_shape)
-       
+        W_f = get_weights(self.W_shape, "v_W", mask=self.mask)
+        b_f = get_bias(self.b_shape, "v_b")
+        W_g = get_weights(self.W_shape, "h_W", mask=self.mask)
+        b_g = get_bias(self.b_shape, "h_b")
+      
         conv_f = conv_op(self.fan_in, W_f)
         conv_g = conv_op(self.fan_in, W_g)
        
@@ -55,8 +54,8 @@ def gated_conv(self):
         self.fan_out = tf.mul(tf.tanh(conv_f + b_f), tf.sigmoid(conv_g + b_g))
 
     def simple_conv(self):
-        W = get_weights(self.W_shape, mask=self.mask)
-        b = get_bias(self.b_shape)
+        W = get_weights(self.W_shape, "W", mask=self.mask)
+        b = get_bias(self.b_shape, "b")
         conv = conv_op(self.fan_in, W)
         if self.activation: 
             self.fan_out = tf.nn.relu(tf.add(conv, b))
diff --git a/pixelcnn_tf.py b/pixelcnn_tf.py
index 5feb336..b94c5b5 100644
--- a/pixelcnn_tf.py
+++ b/pixelcnn_tf.py
@@ -1,10 +1,10 @@
 # TODO
 # try changing 0.0 to np.random
-# make for imagenet data
 # check network arch
 # upscale-downscale-q_level
-# cost on test set
 # autoencoder
+# cost on test set
+# make for imagenet data
 
 import tensorflow as tf
 import numpy as np
@@ -38,7 +38,7 @@
     i = str(i)
 
     with tf.variable_scope("v_stack"+i):
-        v_stack = PixelCNN([FILTER_SIZE, FILTER_SIZE, in_dim, F_MAP], [F_MAP], v_stack_in, mask=mask).output()
+        v_stack = PixelCNN([FILTER_SIZE, FILTER_SIZE, in_dim, F_MAP], [F_MAP], v_stack_in, mask='a').output()
         v_stack_in = v_stack
 
     with tf.variable_scope("v_stack_1"+i):
@@ -97,8 +97,8 @@ def generate_and_save(sess):
                     .reshape([batch_size, img_height, img_width, 1]))
             _, cost = sess.run([optimizer, loss], feed_dict={X:batch_X})
 
-        print "Epoch: %d, Cost: %.2f"%(i, cost)
+        print "Epoch: %d, Cost: %f"%(i, cost)
         generate_and_save(sess)
 
-generate_and_save(sess)
+    generate_and_save(sess)