From e8c1f7a98ebb18e6136f2a22d618caac4aab39fa Mon Sep 17 00:00:00 2001 From: Jeffrey Lai Date: Sun, 3 Aug 2025 22:35:28 -0500 Subject: [PATCH] nit: fix axis in quantile loss computation --- src/chronos/chronos_bolt.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/chronos/chronos_bolt.py b/src/chronos/chronos_bolt.py index f099e04..8ad3172 100644 --- a/src/chronos/chronos_bolt.py +++ b/src/chronos/chronos_bolt.py @@ -363,7 +363,7 @@ def forward( ) * target_mask.float() ) - loss = loss.mean(dim=-2) # Mean over prediction horizon + loss = loss.mean(dim=-1) # Mean over prediction horizon loss = loss.sum(dim=-1) # Sum over quantile levels loss = loss.mean() # Mean over batch