Using a new activity function creates -nan values in TensorFlow

I define a new activity function motivated by the reLU activation. However, it does not function correctly and I get -nan values after the first epoch.

import numpy as np
import tensorflow as tf
from tensorflow import keras 
from tensorflow.python.ops import math_ops

training_input=  np.random.random ([200,5]) * 200 #input data
training_input.astype(int)

training_output = np.random.random ([200,1]) # output data


def custom_act(alpha=0.0):  #new activity function      
    def new_act(x):
        neg = tf.cast(tf.math.greater_equal(0.0, x), dtype='float32')  # less than zero
        pos = tf.cast(tf.math.greater_equal(x, 0.0), dtype='float32') # greater than zero
        # if <0, then x. if>0, then scale it with alpha and exponential
        return  tf.cast( (neg*x) + (alpha*pos*(1-tf.math.exp(-x)) ) , tf.float32)
    return new_act  

model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(5,)),
    tf.keras.layers.Dense(64, activation="relu"), 
    tf.keras.layers.Dense(32,activation=custom_act(alpha=0.1)),
    tf.keras.layers.Dense( 1) 
]) 

model.compile(loss="mse",optimizer =tf.keras.optimizers.Adam(learning_rate=0.1))

model.fit(training_input,training_output,epochs = 5,batch_size=20,verbose=1)

For instance, the first two columns in the first row are -16.0825386 43.0274544 and they turn into -16.0825386 0.1. After the first epoch, then I get all -nan. What might I be doing wrong?

[-nan -nan -nan ... -nan -nan -nan]
[-nan -nan -nan ... -nan -nan -nan]

from Recent Questions - Stack Overflow https://ift.tt/3iiaCtc
https://ift.tt/eA8V8J