diff --git a/climax/main/downscale_train_season.py b/climax/main/downscale_train_season.py
index 20bfc191880f059a29728379b79ec98ce973be5a..ccac375142bd8131e3141dd2276d9b76b6b1084e 100644
--- a/climax/main/downscale_train_season.py
+++ b/climax/main/downscale_train_season.py
@@ -160,7 +160,7 @@ if __name__ == '__main__':
 
         # define number of output fields
         # check whether modelling pr with probabilistic approach
-        outputs = len(Obs_ds.train)
+        outputs = len(Obs_train.data_vars)
         if PREDICTAND == 'pr' and (isinstance(LOSS, BernoulliGammaLoss) or
                                    isinstance(LOSS, BernoulliGenParetoLoss)):
             outputs = 3