i```
nput_shape = (28,28.1)
class cnn_model(tf.keras.Model):
    def __init__(self,inputs=(28,28,1)):
        super(cnn_model,self).__init__()
        
        #self.conv1 = layers.Conv2D(32,(3,3),activation='relu',input_shape= input_shape)
        self.conv1 = layers.Conv2D(32, 3, 3, padding='same', activation='relu')
        self.maxpool = layers.MaxPool2D((2,2))
        self.conv2 = layers.Conv2D(64,(3,3),activation ='relu')
        self.conv3 = layers.Conv2D(128,(3,3),activation='relu')
        self.flatten = layers.Flatten()
        self.dense64 = layers.Dense(64,activation='relu')
        self.dense10 = layers.Dense(10,activation='relu')
        self.dropout = layers.Dropout(0.25)
    def call(self,x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.maxpool(x)
        x = self.dropout(x)
        x = self.flatten(x)
        x = self.dense64(x)
        x = self.dense10(x)
        return x

#loss = tf.losses.mean_squared_error(labels,logits)
cnn_optimizer = tf.optimizers.Adam(1e-4)
#optimizer = tf.train. (learning_rate=0.001)
#print(loss)


checkpoint_dir = './training_checkpoints'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(cnn_optimizer=cnn_optimizer)

# Notice the use of `tf.function`
# This annotation causes the function to be "compiled".
cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)

#@tf.function
def train_step(logits,labels):    
    
    with tf.GradientTape() as cnn_tape:
        print("train_step")
        loss = tf.losses.mean_squared_error(labels,tf.argmax(logits,axis=1))
        #loss = tf.losses.BinaryCrossentropy(multi_class_labels=labels, logits=logits)
        #loss = cross_entropy(labels,tf.argmax(logits,axis=1))
        gradient_of_cnn = cnn_tape.gradient(loss,model.trainable_variables)
        cnn_optimizer.apply_gradients(zip(gradient_of_cnn,model.trainable_variables))

def train(train_dataset,train_labels_dataset,epochs):
    for epoch in range(epochs):
        start = time.time()
        
        for train_batch,label_batch in zip(train_dataset,train_labels_dataset):
            print(train_batch.shape,label_batch.shape)
            #logits = tf.argmax(model(train_batch),axis=1)
            logits = model(train_batch)
            labels = label_batch
            print('logits->',logits.shape , 'labels->',labels.shape)
            train_step(logits,labels)
            
        if (epoch + 1) % 15 == 0:
            checkpoint.save(file_prefix = checkpoint_prefix)
        print ('Time for epoch {} is {} sec'.format(epoch + 1, time.time()-start))

%%time
train(train_dataset,train_labels_dataset, EPOCHS)

tensorflow/lib/python2.7/site-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.
  from ._conv import register_converters as _register_converters
Using TensorFlow backend.
('X_train shape:', (60000, 28, 28, 1))
(60000, 'train samples')
(10000, 'test samples')
Loaded MNIST test data.
Traceback (most recent call last):
  File "/anaconda2/envs/tensorflow/lib/python2.7/runpy.py", line 174, in _run_module_as_main
    "__main__", fname, loader, pkg_name)
  File "/anaconda2/envs/tensorflow/lib/python2.7/runpy.py", line 72, in _run_code
    exec code in run_globals
  File "/DeepLearning/ensemble-adv-training/train.py", line 57, in <module>
    main(args.model, args.type)
  File "/DeepLearning/ensemble-adv-training/train.py", line 39, in main
    tf_train(x, y, model, X_train, Y_train, data_gen)
  File "tf_utils.py", line 79, in tf_train
    optimizer = tf.train.AdamOptimizer().minimize(loss)
  File "/anaconda2/envs/tensorflow/lib/python2.7/site-packages/tensorflow/python/training/optimizer.py", line 350, in minimize
    ([str(v) for _, v in grads_and_vars], loss))
ValueError: No gradients provided for any variable, check your graph for ops that do not support gradients, between variables ["<tf.Variable 'conv2d_1/kernel:0' shape=(5, 5, 1, 64) dtype=float32_ref>", "<tf.Variable 'conv2d_1/bias:0' shape=(64,) dtype=float32_ref>", "<tf.Variable 'conv2d_2/kernel:0' shape=(5, 5, 64, 64) dtype=float32_ref>", "<tf.Variable 'conv2d_2/bias:0' shape=(64,) dtype=float32_ref>", "<tf.Variable 'dense_1/kernel:0' shape=(25600, 128) dtype=float32_ref>", "<tf.Variable 'dense_1/bias:0' shape=(128,) dtype=float32_ref>", "<tf.Variable 'dense_2/kernel:0' shape=(128, 10) dtype=float32_ref>", "<tf.Variable 'dense_2/bias:0' shape=(10,) dtype=float32_ref>"] and loss Tensor("Mean:0", shape=(), dtype=float32).