The formula of adaptive_std in your code may be wrong

[josephlau]

fn = lambda elem: K.switch(K.less_equal(elem, 1.0), K.ones_like(elem), elem)
adaptive_std = K.map_fn(fn, adaptive_std)

In author's code is
adaptive_std[adaptive_std <= self.eps] = 1

which self.eps = 1e-8

[josephlau]

Btw, change to tf.where() from map_fn results much faster training time

[llouislu]

@josephlau Can you please test my updated code https://github.com/llouislu/tensorflow-dain/tree/bugfix-1 if you're still interested? as I don't have access to data and environment to test dain at the moment.