NaN occurred in optimization at first Iteration with ADVI

The model looks fine to me - maybe try changing the TruncatedNormal to HalfNormal?