class ResidualBlock(Layer):
def __init__(self, **kwargs):
    super(ResidualBlock, self).__init__(**kwargs)
    
def build(self, input_shape):
    """
    This method should build the layers according to the above specification. Make sure 
    to use the input_shape argument to get the correct number of filters, and to set the
    input_shape of the first layer in the block.
    """
    self.batch_norm1 = BatchNormalization(input_shape=input_shape)
    self.conv1 = Conv2D(filters=input_shape[-1], kernel_size=(3,3), padding='SAME')
    self.batch_norm2 = BatchNormalization()
    self.conv2 = Conv2D(filters=input_shape[-1], kernel_size=(3,3), padding='SAME')
    
def call(self, inputs, training=False):
    """
    This method should contain the code for calling the layer according to the above
    specification, using the layer objects set up in the build method.
    """
    x = self.batch_norm1(inputs, training=training)
    x = tf.nn.relu(x)
    x = self.conv1(x)
    x = self.batch_norm2(x, training=training)
    x = tf.nn.relu(x)
    outputs = self.conv2(x)
    
    return tf.add(outputs, inputs)

class FiltersChangeResidualBlock(Layer):
def __init__(self, out_filters, **kwargs):
    """
    The class initialiser should call the base class initialiser, passing any keyword
    arguments along. It should also set the number of filters as a class attribute.
    """
    super(FiltersChangeResidualBlock, self).__init__(**kwargs)
    self.out_filters = out_filters
    
def build(self, input_shape):
    """
    This method should build the layers according to the above specification. Make sure 
    to use the input_shape argument to get the correct number of filters, and to set the
    input_shape of the first layer in the block.
    """
    self.batch_norm1 = BatchNormalization(input_shape=input_shape)
    self.conv1 = Conv2D(filters=input_shape[-1], kernel_size=(3,3), padding='SAME')
    self.batch_norm2 = BatchNormalization()
    self.conv2 = Conv2D(filters=self.out_filters, kernel_size=(3,3), padding='SAME')
    self.conv3 = Conv2D(filters=self.out_filters, kernel_size=(1,1))
    
def call(self, inputs, training=False):
    """
    This method should contain the code for calling the layer according to the above
    specification, using the layer objects set up in the build method.
    """
    x = self.batch_norm1(inputs, training=training)
    x = tf.nn.relu(x)
    x = self.conv1(x)
    x = self.batch_norm2(x, training=training)
    x = tf.nn.relu(x)
    x = self.conv2(x)
    return tf.add(x, self.conv3(inputs))

class ResNetModel(Model):
def __init__(self, **kwargs):
    """
    The class initialiser should call the base class initialiser, passing any keyword
    arguments along. It should also create the layers of the network according to the
    above specification.
    """
    super(ResNetModel, self).__init__(**kwargs)
    self.conv1 = Conv2D(32, (7,7), strides= 2)
    self.residual1 = ResidualBlock()
    self.conv2 = Conv2D(32, (3,3), strides=2)
    self.residual2 = FiltersChangeResidualBlock(64)
    self.flatten = Flatten()
    self.dense = Dense(10, activation='softmax')
    
def call(self, inputs, training=False):
    """
    This method should contain the code for calling the layer according to the above
    specification, using the layer objects set up in the initialiser.
    """
    x = self.conv1(inputs)
    x = self.residual1(x, training=training)
    x = self.conv2(x)
    x = self.residual2(x, training=training)
    x = self.flatten(x)
    output = self.dense(x)
    
    return output

@tf.function
def grad(model, inputs, targets, loss):
“”"
This function should compute the loss and gradients of your model, corresponding to
the inputs and targets provided. It should return the loss and gradients.
“”"
with tf.GradientTape() as tape:
loss_value = loss(model, inputs, targets)
return (loss_value, tape.gradient(loss_values, model.trainable_variables))

def train_resnet(model, num_epochs, dataset, optimizer, loss, grad_fn):
“”"
This function should implement the custom training loop, as described above. It should
return a tuple of two elements: the first element is a list of loss values per epoch, the
second is a list of accuracy values per epoch
“”"
#train_dataset = tf.data.Dataset.from_tensor_slices((x_train, train_labels))
#train_dataset = dataset.batch(32)
train_dataset = dataset 

train_loss_results = []
train_accuracy_results = []

#num_epochs = 10
#weight_decay = .005

for epoch in range(num_epochs):

    epoch_loss_avg = tf.keras.metrics.Mean()
    epoch_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
    print(epoch)
    for x, y in train_dataset:
        #print(x, y)
        loss_value, grads = grad_fn(model, x, y, loss)
        print(grads)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))

        #Compute Current loss:
        epoch_loss_avg(loss_value)
        #Compare predicted label to actual label:
        epoch_accuracy(to_categorical(y), model(x))


    #End Epoch:
    train_loss_results.append(epoch_loss_avg.result())
    train_accuracy_results.append(epoch_accuracy.result())

    print('Epoch {} Loss {:.3f} Accuracy {:.3f}'.format(epoch+1, 
                                                       epoch_loss_avg.result(),
                                                       epoch_accuracy.result()))

return train_loss_results, train_accuracy_results`