Warnings with so many “..xx/while/xx..” in it pop up as model's size increased

Question

I try to train tfd.TransformedDistribution with stack of tfb.ffjord as bijectors. Basically, more bijectors make model deeper and having more param