def test_tf2_training(art_warning, fix_get_mnist_data, fix_get_cifar10_data):
    """
    Check that the training loop for tensorflow2 does not result in errors
    """
    import tensorflow as tf

    def build_model(input_shape):
        img_inputs = tf.keras.Input(shape=(input_shape[0], input_shape[1], input_shape[2] * 2))
        x = tf.keras.layers.Conv2D(filters=32, kernel_size=(4, 4), strides=(2, 2), activation="relu")(img_inputs)
        x = tf.keras.layers.MaxPool2D(pool_size=(2, 2), strides=2)(x)
        # tensorflow uses channels last, and we are loading weights from an originally trained pytorch model
        x = tf.transpose(x, (0, 3, 1, 2))
        x = tf.keras.layers.Flatten()(x)
        x = tf.keras.layers.Dense(100, activation="relu")(x)
        x = tf.keras.layers.Dense(10)(x)
        return tf.keras.Model(inputs=img_inputs, outputs=x)

    loss_object = tf.keras.losses.CategoricalCrossentropy(from_logits=True)
    optimizer = tf.keras.optimizers.legacy.SGD(learning_rate=0.01)

    for dataset, dataset_name in zip([fix_get_mnist_data, fix_get_cifar10_data], ["mnist", "cifar"]):
        if dataset_name == "mnist":
            input_shape = (28, 28, 1)
        else:
            input_shape = (32, 32, 3)
        net = build_model(input_shape=input_shape)

        try:
            for ablation_type in ["column", "row", "block"]:
                ablation_size = 5
                classifier = TensorFlowV2DeRandomizedSmoothing(
                    model=net,
                    clip_values=(0, 1),
                    loss_object=loss_object,
                    optimizer=optimizer,
                    input_shape=input_shape,
                    nb_classes=10,
                    ablation_type=ablation_type,
                    ablation_size=ablation_size,
                    threshold=0.3,
                    logits=True,
                )
                x = np.transpose(np.copy(dataset[0]), (0, 2, 3, 1))  # put channels last
                classifier.fit(x=x, y=dataset[1], nb_epochs=1)
        except ARTTestException as e:
            art_warning(e)

    Check that the training loop for tensorflow2 does not result in errors
    
tests/estimators/certification/test_derandomized_smoothing.py