One more point. Today I update the parameter initialization, changing the bias initialization from 1 to 0. Now I can initialize the denoising params to 0 an everything works. It seems to be an internal problem in pytorch, as the input to the derivable variable is:
mu_i=a1+a2·z_proj·a3+a4+a5
v_i=a6+a7·z_proj·a8+a9+a10
where in the upper layer the proj comes from a batch normalization of the sotmax activation of a linear projection +BN+ noise +add batch norma parameters. The bias is included as the beta of the batchnorm, this means, the linear projection only perform matrix multiplication.
Then just use
z_denoised = (z_level-mu_i)*v_i+mu_i
where z_level is the linear projection +BN+ noise
This should not give nan for having a bias initialize to one…, when evaluating the MSE of the z_denoised vs the z_clean, which have same transformation as the z_level except that we should not add noise to it.
Update:
I change some things  in the code. The point was not initialize the bias to 0. It has to do with batch normalization. When finish I will push the code to github so the pytorch developers can fix the problem.