Fast gradient sign method with keras

I'm currently working on this paper . To implement the Fast gradient sign method with a heteroscedastic neural network.

If we define the loss function as l(\theta,x,y) where x is the feature, y the label and \theta the parameters. Instead of minimizing l(\theta,x,y), the goal is to minimize l(\theta,x,y)+l(\theta,x',y) where

x'=x+\eps*\sign(\nabla_x l(\theta,x,y))

Here my attempt(without any success):

def customLoss(x):
    def neg_log_likelihood(y_true, y_pred):

        def neg_log(y_t,y_p):
            inter=(y_p[...,0,None]-y_t)/K.clip(y_p[...,1,None],K.epsilon(),None)
            val=K.log(K.clip(K.square(y_p[...,1,None]),K.epsilon(),None))+K.square(inter)
            return val

        val=neg_log(y_true,y_pred)

        deriv=K.gradients(val,x)
        xb=x+0.01*K.sign(deriv)
        out=model.predict(xb)
        valb=neg_log(y_true,out)

        return K.mean(val+valb,axis=-1)
    return neg_log_likelihood

The calling of this loss function is then

model.compile(loss=customLoss(model.inputs),...)

Do you have any ideas how can I implement this?