Fix device placement of InverseGamma distribution
@@ -23,7 +23,8 @@ class InverseGamma(TransformedDistribution):
def __init__(self, concentration, rate, validate_args=None):
base_dist = Gamma(concentration, rate)
- super().__init__(base_dist, PowerTransform(-1.0), validate_args=validate_args)
+ super().__init__(base_dist, PowerTransform(-base_dist.rate.new_ones(())),
+ validate_args=validate_args)
def expand(self, batch_shape, _instance=None):
new = self._get_checked_instance(InverseGamma, _instance)
