From 68a850741c65e37b7ec5366a5e98d0cf07dde0c2 Mon Sep 17 00:00:00 2001 From: t-minus Date: Sun, 16 Mar 2025 23:05:45 +0000 Subject: [PATCH] Issue 712 - fix scaling for arguments --- nbs/losses.pytorch.ipynb | 4 ++-- neuralforecast/losses/pytorch.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) 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)