simple_conv2.py

print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))

#input image is 5 X 5 and 1 channel 
input_shape = (1, 1, 5, 5)
tf.summary.trace_on(graph=True, profiler=True)

x = tf.random.normal(input_shape) 
print(x.shape)

y = tf.keras.layers.Conv2D( 
2, 2, activation='relu', input_shape=input_shape, data_format='channels_first')(x) 
print(y.shape)

print(x)

#output 2 feature map 5 X 5 each 
print(y)