Convert tensorflow_probability.distributions.TransformedDistribution to pytorch

1

I am trying to convert a tensorflow script to pytorch. I will appreciate if someone can explain what would be the equivalent piece of code for the following line from this script

import tensorflow_probability as tfp
tfd = tfp.distributions
a_distribution = tfd.TransformedDistribution(
        distribution=tfd.Normal(loc=0.0, scale=1.0),
        bijector=tfp.bijectors.Chain([
            tfp.bijectors.AffineScalar(shift=self._action_means,
                                       scale=self._action_mags),
            tfp.bijectors.Tanh(),
            tfp.bijectors.AffineScalar(shift=mean, scale=std),
        ]),
        event_shape=[mean.shape[-1]],
        batch_shape=[mean.shape[0]])

Thanks

2

I converted the above part in pytorch as following

a_distribution = TransformedDistribution(
                        base_distribution=Normal(loc=torch.full_like(mean, 0), 
                                                 scale=torch.full_like(mean, 1)), 
                        transforms=tT.ComposeTransform([
                                   tT.AffineTransform(loc=self._action_means, scale=self._action_mags, event_dim=mean.shape[-1]), 
                                   tT.TanhTransform(),
                                   tT.AffineTransform(loc=mean, scale=std, event_dim=mean.shape[-1])]))

But TanhTransform caused a nan value. I am wondering whether I have done this incorrectly or not?