making a picture of a neural network in tensorflow

Question

I'm new to Machine Learning and I have just gone through a tutorial explaining how to create a neural network in TensorFlow. I was wondering if it is possible to visualize the neural network I created. The output should be a picture like this

Thanks.

MRE:

import tensorflow as tf

ann = tf.keras.models.Sequential()
ann.add(tf.keras.layers.Dense(units=6, activation='relu'))
ann.add(tf.keras.layers.Dense(units=6, activation='relu'))
ann.add(tf.keras.layers.Dense(units=1, activation='sigmoid'))
tf.keras.utils.plot_model(ann, to_file='model.png', show_shapes=True, 
                          show_layer_names=True, rankdir='TB',
                          expand_nested=False, dpi=96
                          )

Answer 1

This should do the trick:

tf.keras.utils.plot_model(
        model, to_file='model.png', show_shapes=False, show_layer_names=True,
        rankdir='TB', expand_nested=False, dpi=96
    )

This will generate an image like:

Example: https://www.tensorflow.org/guide/keras/functional