network.py

import tensorflow as tf

def build_nn(input_shape, hidden_units=(20, 20), output_size=1, output_activation=None):
    inputs = tf.keras.Input(shape=input_shape)
    
    hidden = inputs
    for layer_units in hidden_units:
        hidden = tf.keras.layers.Dense(layer_units, activation=tf.nn.relu)(hidden)
        
    outputs = tf.keras.layers.Dense(output_size, activation=output_activation)(hidden)
    
    model = tf.keras.Model(inputs=inputs, outputs=outputs)
    return model