import tensorflow as tf

            if optimizer == 'sgd':
                lr = tf.train.exponential_decay(
                    learning_rate,
                    global_step,
                    1,
                    learning_rate_decay
                )
                train_step = tf.train.GradientDescentOptimizer(lr).minimize(loss, global_step=global_step)
            elif optimizer == 'adam':
                train_step = tf.train.GradientDescentOptimizer(learning_rate).\
                                                               minimize(loss, global_step=global_step)     
            else:
                raise NotImplementedError('activation not recognized')   
            # init op
            init_op = tf.global_variables_initializer()
        
        # Save into class members
        self.X = X
