diff --git a/nbs/losses.pytorch.ipynb b/nbs/losses.pytorch.ipynb index d3ff48894..65dfc7605 100644 --- a/nbs/losses.pytorch.ipynb +++ b/nbs/losses.pytorch.ipynb @@ -1454,8 +1454,8 @@ " mu = F.softplus(mu) + 1e-8\n", " alpha = F.softplus(alpha) + 1e-8 # alpha = 1/total_counts\n", " if (loc is not None) and (scale is not None):\n", - " mu *= loc\n", - " alpha /= (loc + 1.)\n", + " mu = mu * scale + loc\n", + " alpha /= (scale + 1.)\n", "\n", " # mu = total_count * (probs/(1-probs))\n", " # => probs = mu / (total_count + mu)\n", diff --git a/neuralforecast/losses/pytorch.py b/neuralforecast/losses/pytorch.py index f24d2e5bc..03b1a1345 100644 --- a/neuralforecast/losses/pytorch.py +++ b/neuralforecast/losses/pytorch.py @@ -857,8 +857,8 @@ def nbinomial_scale_decouple(output, loc=None, scale=None): mu = F.softplus(mu) + 1e-8 alpha = F.softplus(alpha) + 1e-8 # alpha = 1/total_counts if (loc is not None) and (scale is not None): - mu *= loc - alpha /= loc + 1.0 + mu = mu * scale + loc + alpha /= scale + 1.0 # mu = total_count * (probs/(1-probs)) # => probs = mu / (total_count + mu)