diff --git a/climax/main/config.py b/climax/main/config.py index 56fcbceefabb7e7c531fc70bcb017c3e2ff7f144..a6a97a6370251a41ee372282776ef5b8fd13a355 100644 --- a/climax/main/config.py +++ b/climax/main/config.py @@ -149,7 +149,11 @@ if PREDICTAND is 'pr': if isinstance(LOSS, MSELoss): MAX_LR = 0.0004 if isinstance(LOSS, BernoulliGammaLoss): - MAX_LR = 0.0005 if OPTIM is torch.optim.Adam else 0.001 + # learning rates for supersampling task + if not any(ERA5_P_PREDICTORS) and ERA5_S_PREDICTORS == 'pr': + MAX_LR = 0.002 + else: + MAX_LR = 0.0005 if OPTIM is torch.optim.Adam else 0.001 # base learning rate: MAX_LR / 4 (Smith L. (2017)) BASE_LR = MAX_LR / 4