autogluon.tabular 1.5.1b20260105__py3-none-any.whl → 1.5.1b20260116__py3-none-any.whl

autogluon/tabular/models/mitra/_internal/config/config_pretrain.py

@@ -1,18 +1,19 @@
 from __future__ import annotations
 
+import os
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional
-import yaml
-import os
 
 import torch
+import yaml
 from omegaconf import DictConfig, OmegaConf
 
-from ..._internal.config.enums import GeneratorName, ModelName, LossName, Task
+from ..._internal.config.enums import GeneratorName, LossName, ModelName, Task
+
 
 @dataclass
-class ConfigData():
+class ConfigData:
     generator: GeneratorName
     min_samples_support: int
     max_samples_support: int
@@ -26,37 +27,38 @@ class ConfigData():
     task: Task
 
     def __post_init__(self):
-
         assert self.min_samples_support <= self.max_samples_support
         assert self.min_features <= self.max_features
 
+
 @dataclass
-class ConfigModel():
+class ConfigModel:
     name: ModelName
     hyperparams: dict
 
 
 @dataclass
-class ConfigPreprocessing():
+class ConfigPreprocessing:
     use_quantile_transformer: bool
     use_feature_count_scaling: bool
 
+
 @dataclass
-class ConfigGradScaler():
+class ConfigGradScaler:
     enabled: bool
     scale_init: float
     scale_min: float
     growth_interval: int
 
-
     def __post_init__(self):
         assert self.scale_init >= self.scale_min, "Scale init must be greater than scale min"
         assert self.scale_min >= 1, "Scale min lower than 1 makes no sense for mixed precision training"
         assert type(self.scale_init) == float, "Scale init must be a float, otherwise gradscaler will return an error"
         assert type(self.scale_min) == float, "Scale min must be a float, otherwise gradscaler will return an error"
 
+
 @dataclass
-class ConfigOptim():
+class ConfigOptim:
     steps: int
     log_every_n_steps: int
     eval_every_n_steps: int
@@ -80,42 +82,35 @@ class ConfigOptim():
 
     @classmethod
     def from_hydra(cls, cfg_hydra: DictConfig) -> Self:
-
         grad_scaler = ConfigGradScaler(**cfg_hydra.grad_scaler)
-        cfg_dict: dict = OmegaConf.to_container(cfg_hydra)      # type: ignore
+        cfg_dict: dict = OmegaConf.to_container(cfg_hydra)  # type: ignore
         del cfg_dict["grad_scaler"]
 
         regression_loss = LossName[cfg_dict["regression_loss"]]
         del cfg_dict["regression_loss"]
 
-        return cls(
-            grad_scaler=grad_scaler,
-            regression_loss=regression_loss,
-            **cfg_dict
-        )
+        return cls(grad_scaler=grad_scaler, regression_loss=regression_loss, **cfg_dict)
 
     def __post_init__(self):
         assert hasattr(torch, self.precision), f"Precision {self.precision} not supported by torch"
 
-class ConfigSaveLoadMixin(yaml.YAMLObject):
 
+class ConfigSaveLoadMixin(yaml.YAMLObject):
     def save(self, path: Path) -> None:
-
         path.parent.mkdir(parents=True, exist_ok=True)
 
-        with open(path, 'w') as f:
+        with open(path, "w") as f:
             yaml.dump(self, f, default_flow_style=False)
 
-
     @classmethod
     def load(cls, path: Path) -> Self:
-
-        with open(path, 'r') as f:
+        with open(path, "r") as f:
             # It's unsafe, but not unsafer than the pickle module
             config = yaml.unsafe_load(f)
 
         return config
 
+
 @dataclass
 class ConfigPretrain(ConfigSaveLoadMixin):
     run_name: str
@@ -141,8 +136,9 @@ class ConfigPretrain(ConfigSaveLoadMixin):
 
     @classmethod
     def from_hydra(cls, cfg_hydra: DictConfig):
-
-        assert not os.path.exists(cfg_hydra.output_dir), f'Output directory {cfg_hydra.output_dir} already exists! Please change to a new folder.'
+        assert not os.path.exists(cfg_hydra.output_dir), (
+            f"Output directory {cfg_hydra.output_dir} already exists! Please change to a new folder."
+        )
 
         output_dir = Path(cfg_hydra.output_dir)
 
@@ -160,11 +156,11 @@ class ConfigPretrain(ConfigSaveLoadMixin):
             use_ddp=len(devices) > 1,
             seed=cfg_hydra.seed,
             workers_per_gpu=cfg_hydra.workers_per_gpu,
-            model = ConfigModel(
-                name = ModelName[cfg_hydra.model.name],
-                hyperparams = OmegaConf.to_container(cfg_hydra.model.hyperparams),   
+            model=ConfigModel(
+                name=ModelName[cfg_hydra.model.name],
+                hyperparams=OmegaConf.to_container(cfg_hydra.model.hyperparams),
             ),
-            data = ConfigData(
+            data=ConfigData(
                 generator=GeneratorName(cfg_hydra.data.generator),
                 min_samples_support=cfg_hydra.data.min_samples_support,
                 max_samples_support=cfg_hydra.data.max_samples_support,
@@ -175,16 +171,16 @@ class ConfigPretrain(ConfigSaveLoadMixin):
                 task=Task[cfg_hydra.data.task],
                 sample_multinomial_categorical=cfg_hydra.data.sample_multinomial_categorical,
                 sample_multinomial_label=cfg_hydra.data.sample_multinomial_label,
-                generator_hyperparams=OmegaConf.to_container(cfg_hydra.data.generator_hyperparams),    # type: ignore
+                generator_hyperparams=OmegaConf.to_container(cfg_hydra.data.generator_hyperparams),  # type: ignore
             ),
-            optim = ConfigOptim.from_hydra(cfg_hydra.optim),
-            preprocessing = ConfigPreprocessing(**cfg_hydra.preprocessing),
-            load_from_file = cfg_hydra.load_from_file,
-            load_path_x = cfg_hydra.load_path_x,
-            load_path_y = cfg_hydra.load_path_y,
-            save_file = cfg_hydra.save_file,
-            save_file_only = cfg_hydra.save_file_only,
-            save_path_x = cfg_hydra.save_path_x,
-            save_path_y = cfg_hydra.save_path_y,
-            number_of_runs = cfg_hydra.number_of_runs,
+            optim=ConfigOptim.from_hydra(cfg_hydra.optim),
+            preprocessing=ConfigPreprocessing(**cfg_hydra.preprocessing),
+            load_from_file=cfg_hydra.load_from_file,
+            load_path_x=cfg_hydra.load_path_x,
+            load_path_y=cfg_hydra.load_path_y,
+            save_file=cfg_hydra.save_file,
+            save_file_only=cfg_hydra.save_file_only,
+            save_path_x=cfg_hydra.save_path_x,
+            save_path_y=cfg_hydra.save_path_y,
+            number_of_runs=cfg_hydra.number_of_runs,
         )

autogluon/tabular/models/mitra/_internal/config/config_run.py

@@ -16,17 +16,5 @@ class ConfigRun(ConfigSaveLoadMixin):
     hyperparams: dict
 
     @classmethod
-    def create(
-        cls,  
-        device: torch.device, 
-        seed: int, 
-        model_name: ModelName, 
-        hyperparams: dict
-    ) -> "ConfigRun":
-       
-        return cls(
-            device=device,
-            seed=seed,
-            model_name=model_name,
-            hyperparams=hyperparams
-        )
+    def create(cls, device: torch.device, seed: int, model_name: ModelName, hyperparams: dict) -> "ConfigRun":
+        return cls(device=device, seed=seed, model_name=model_name, hyperparams=hyperparams)

autogluon/tabular/models/mitra/_internal/config/enums.py

@@ -5,10 +5,12 @@ try:
 except ImportError:
     # StrEnum is not available in Python < 3.11, so we create a compatible version
     from enum import Enum
+
     class StrEnum(str, Enum):
         """
         Enum where members are also (and must be) strings
         """
+
         def __new__(cls, value):
             if not isinstance(value, str):
                 raise TypeError(f"{value!r} is not a string")
@@ -53,7 +55,7 @@ class Phase(StrEnum):
 
 
 class ModelName(StrEnum):
-    PLACEHOLDER = "_placeholder_"   # This is a placeholder for the current running model
+    PLACEHOLDER = "_placeholder_"  # This is a placeholder for the current running model
     FT_TRANSFORMER = "FT-Transformer"
     TABPFN = "TabPFN"
     FOUNDATION = "Foundation"
@@ -87,10 +89,10 @@ class ModelName(StrEnum):
 
 
 class ModelClass(StrEnum):
-    BASE = 'base'
-    GBDT = 'GBDT'
-    NN = 'NN'
-    ICLT = 'ICLT'
+    BASE = "base"
+    GBDT = "GBDT"
+    NN = "NN"
+    ICLT = "ICLT"
 
 
 class DownstreamTask(StrEnum):
@@ -98,7 +100,6 @@ class DownstreamTask(StrEnum):
     FINETUNE = "finetune"
 
 
-
 class BenchmarkName(StrEnum):
     DEBUG_CLASSIFICATION = "debug_classification"
     DEBUG_REGRESSION = "debug_regression"
@@ -124,25 +125,25 @@ class BenchmarkOrigin(StrEnum):
 
 
 class GeneratorName(StrEnum):
-    TABPFN = 'tabpfn'
-    TREE = 'tree'
-    RANDOMFOREST = 'randomforest'
-    NEIGHBOR = 'neighbor'
-    MIX = 'mix'
-    PERLIN = 'perlin'
-    MIX_7 = 'mix_7'
-    MIX_6 = 'mix_6'
-    MIX_5 = 'mix_5'
-    MIX_5_GP = 'mix_5_gp'
-    MIX_4 = 'mix_4'
-    MIX_4_AG = 'mix_4_ag'
-    LR = 'lr'
-    POLY = 'poly'
-    SAMPLE_RF = 'sample_rf'
-    SAMPLE_GP = 'sample_gp'
-    TABREPO = 'tabrepo'
-    MIX_4_TABREPO = 'mix_4_tabrepo'
-    MIX_4_TABPFNV2 = 'mix_4_tabpfnv2'
+    TABPFN = "tabpfn"
+    TREE = "tree"
+    RANDOMFOREST = "randomforest"
+    NEIGHBOR = "neighbor"
+    MIX = "mix"
+    PERLIN = "perlin"
+    MIX_7 = "mix_7"
+    MIX_6 = "mix_6"
+    MIX_5 = "mix_5"
+    MIX_5_GP = "mix_5_gp"
+    MIX_4 = "mix_4"
+    MIX_4_AG = "mix_4_ag"
+    LR = "lr"
+    POLY = "poly"
+    SAMPLE_RF = "sample_rf"
+    SAMPLE_GP = "sample_gp"
+    TABREPO = "tabrepo"
+    MIX_4_TABREPO = "mix_4_tabrepo"
+    MIX_4_TABPFNV2 = "mix_4_tabpfnv2"
 
 
 class MetricName(StrEnum):
@@ -159,4 +160,4 @@ class MetricName(StrEnum):
 class LossName(StrEnum):
     CROSS_ENTROPY = "cross_entropy"
     MSE = "mse"
-    MAE = "mae"
+    MAE = "mae"

autogluon/tabular/models/mitra/_internal/core/__init__.py

@@ -1 +1 @@
-# Core modules for MitraModel 
+# Core modules for MitraModel

autogluon/tabular/models/mitra/_internal/core/callbacks.py

@@ -2,10 +2,8 @@ import numpy as np
 import torch
 
 
-class EarlyStopping():
-
-    def __init__(self, patience=10, delta=0.0001, metric='log_loss'):
-
+class EarlyStopping:
+    def __init__(self, patience=10, delta=0.0001, metric="log_loss"):
         self.patience = patience
         self.counter = 0
         self.best_score = None
@@ -13,9 +11,7 @@ class EarlyStopping():
         self.delta = delta
         self.metric = metric
 
-
     def __call__(self, val_loss):
-        
         # smaller is better for these metrics
         if self.metric in ["log_loss", "mse", "mae", "rmse"]:
             score = -val_loss
@@ -23,7 +19,9 @@ class EarlyStopping():
         elif self.metric in ["accuracy", "roc_auc", "r2"]:
             score = val_loss
         else:
-            raise ValueError(f"Unsupported metric: {self.metric}. Supported metrics are: log_loss, mse, mae, rmse, accuracy, roc_auc, r2.")
+            raise ValueError(
+                f"Unsupported metric: {self.metric}. Supported metrics are: log_loss, mse, mae, rmse, accuracy, roc_auc, r2."
+            )
 
         if self.best_score is None:
             self.best_score = score
@@ -39,39 +37,34 @@ class EarlyStopping():
         return self.early_stop
 
 
-class Checkpoint():
-
+class Checkpoint:
     def __init__(self):
         self.curr_best_loss = np.inf
         self.best_model: dict
-    
+
     def reset(self, net: torch.nn.Module):
         self.curr_best_loss = np.inf
         self.best_model = net.state_dict()
         for key in self.best_model:
-            self.best_model[key] = self.best_model[key].to('cpu')
-        
+            self.best_model[key] = self.best_model[key].to("cpu")
 
     def __call__(self, net: torch.nn.Module, loss: float):
-        
         if loss < self.curr_best_loss:
             self.curr_best_loss = loss
             self.best_model = net.state_dict()
             for key in self.best_model:
-                self.best_model[key] = self.best_model[key].to('cpu')
-
+                self.best_model[key] = self.best_model[key].to("cpu")
 
     def set_to_best(self, net):
         net.load_state_dict(self.best_model)
 
 
-class EpochStatistics():
-
+class EpochStatistics:
     def __init__(self) -> None:
         self.n = 0
         self.loss = 0
         self.score = 0
-        
+
     def update(self, loss, score, n):
         self.n += n
         self.loss += loss * n
@@ -79,9 +72,9 @@ class EpochStatistics():
 
     def get(self):
         return self.loss / self.n, self.score / self.n
-    
-class TrackOutput():
 
+
+class TrackOutput:
     def __init__(self) -> None:
         self.y_true: list[np.ndarray] = []
         self.y_pred: list[np.ndarray] = []
@@ -91,4 +84,4 @@ class TrackOutput():
         self.y_pred.append(y_pred)
 
     def get(self):
-        return np.concatenate(self.y_true, axis=0), np.concatenate(self.y_pred, axis=0)
+        return np.concatenate(self.y_true, axis=0), np.concatenate(self.y_pred, axis=0)

autogluon/tabular/models/mitra/_internal/core/get_loss.py

@@ -7,13 +7,11 @@ from ..._internal.config.enums import LossName, Task
 
 
 class CrossEntropyLossExtraBatch(torch.nn.Module):
-
     def __init__(self, label_smoothing: float):
         super().__init__()
 
         self.loss = torch.nn.CrossEntropyLoss(label_smoothing=label_smoothing)
 
-
     def forward(self, input, target):
         """
         Input has shape (batch_size, num_samples, num_classes)
@@ -22,26 +20,26 @@ class CrossEntropyLossExtraBatch(torch.nn.Module):
         Compared to the original CrossEntropyLoss, accepts (batch_size, num_samples) as batch
         """
 
-        input = einops.rearrange(input, 'b s c -> (b s) c')
-        target = einops.rearrange(target, 'b s -> (b s)')
+        input = einops.rearrange(input, "b s c -> (b s) c")
+        target = einops.rearrange(target, "b s -> (b s)")
 
         return self.loss(input, target)
 
-def get_loss(cfg: ConfigRun):
 
-    if cfg.task == Task.REGRESSION and cfg.hyperparams['regression_loss'] == LossName.MSE:
+def get_loss(cfg: ConfigRun):
+    if cfg.task == Task.REGRESSION and cfg.hyperparams["regression_loss"] == LossName.MSE:
         return torch.nn.MSELoss()
-    elif cfg.task == Task.REGRESSION and cfg.hyperparams['regression_loss'] == LossName.MAE:
+    elif cfg.task == Task.REGRESSION and cfg.hyperparams["regression_loss"] == LossName.MAE:
         return torch.nn.L1Loss()
-    elif cfg.task == Task.REGRESSION and cfg.hyperparams['regression_loss'] == LossName.CROSS_ENTROPY:
-        return CrossEntropyLossExtraBatch(cfg.hyperparams['label_smoothing'])
+    elif cfg.task == Task.REGRESSION and cfg.hyperparams["regression_loss"] == LossName.CROSS_ENTROPY:
+        return CrossEntropyLossExtraBatch(cfg.hyperparams["label_smoothing"])
     elif cfg.task == Task.CLASSIFICATION:
-        return CrossEntropyLossExtraBatch(cfg.hyperparams['label_smoothing'])
+        return CrossEntropyLossExtraBatch(cfg.hyperparams["label_smoothing"])
     else:
         raise ValueError(f"Unsupported task {cfg.task} and (regression) loss {cfg.hyperparams['regression_loss']}")
 
-def get_loss_pretrain(cfg: ConfigPretrain):
 
+def get_loss_pretrain(cfg: ConfigPretrain):
     if cfg.data.task == Task.REGRESSION and cfg.optim.regression_loss == LossName.MSE:
         return torch.nn.MSELoss()
     elif cfg.data.task == Task.REGRESSION and cfg.optim.regression_loss == LossName.MAE:
@@ -51,4 +49,4 @@ def get_loss_pretrain(cfg: ConfigPretrain):
     elif cfg.data.task == Task.CLASSIFICATION:
         return CrossEntropyLossExtraBatch(cfg.optim.label_smoothing)
     else:
-        raise ValueError(f"Unsupported task {cfg.data.task} and (regression) loss {cfg.optim.regression_loss}")
+        raise ValueError(f"Unsupported task {cfg.data.task} and (regression) loss {cfg.optim.regression_loss}")

autogluon/tabular/models/mitra/_internal/core/get_optimizer.py

@@ -5,44 +5,32 @@ from ..._internal.config.config_pretrain import ConfigPretrain
 
 
 def get_optimizer(hyperparams: dict, model: torch.nn.Module) -> torch.optim.Optimizer:
-
     optimizer: torch.optim.Optimizer
 
-    if hyperparams['optimizer'] == "adam":
+    if hyperparams["optimizer"] == "adam":
         optimizer = Adam(
-            model.parameters(), 
-            lr=hyperparams['lr'],
-            betas=(0.9, 0.999),
-            weight_decay=hyperparams['weight_decay']
+            model.parameters(), lr=hyperparams["lr"], betas=(0.9, 0.999), weight_decay=hyperparams["weight_decay"]
         )
-    elif hyperparams['optimizer'] == "adamw":
+    elif hyperparams["optimizer"] == "adamw":
         optimizer = AdamW(
-            model.parameters(), 
-            lr=hyperparams['lr'],
-            betas=(0.9, 0.999),
-            weight_decay=hyperparams['weight_decay']
-        )
-    elif hyperparams['optimizer'] == "sgd":
-        optimizer = SGD(
-            model.parameters(),
-            lr=hyperparams['lr'],
-            weight_decay=hyperparams['weight_decay']
+            model.parameters(), lr=hyperparams["lr"], betas=(0.9, 0.999), weight_decay=hyperparams["weight_decay"]
         )
+    elif hyperparams["optimizer"] == "sgd":
+        optimizer = SGD(model.parameters(), lr=hyperparams["lr"], weight_decay=hyperparams["weight_decay"])
     else:
         raise ValueError("Optimizer not recognized")
-    
+
     return optimizer
 
 
 def get_optimizer_pretrain(cfg: ConfigPretrain, model: torch.nn.Module) -> torch.optim.Optimizer:
-
     parameters = [(name, param) for name, param in model.named_parameters()]
 
     parameters_with_weight_decay = []
     parameters_without_weight_decay = []
 
     for name, param in parameters:
-        if name.endswith("bias") or 'norm' in name or 'embedding' in name:
+        if name.endswith("bias") or "norm" in name or "embedding" in name:
             parameters_without_weight_decay.append(param)
         else:
             parameters_with_weight_decay.append(param)
@@ -53,26 +41,25 @@ def get_optimizer_pretrain(cfg: ConfigPretrain, model: torch.nn.Module) -> torch
     ]
 
     optimizer = torch.optim.AdamW(
-        optimizer_parameters, 
+        optimizer_parameters,
         lr=cfg.optim.lr,
         betas=(cfg.optim.beta1, cfg.optim.beta2),
-        weight_decay=cfg.optim.weight_decay
+        weight_decay=cfg.optim.weight_decay,
     )
-    
+
     return optimizer
 
 
 class GradScaler(torch.amp.GradScaler):
-
     def __init__(
-        self, 
+        self,
         enabled: bool = True,
-        scale_init: float = 2.**16,
-        scale_min: float = 1.,
+        scale_init: float = 2.0**16,
+        scale_min: float = 1.0,
         growth_interval: int = 2000,
-        device: str = 'cuda'
+        device: str = "cuda",
     ):
-        super().__init__(enabled=enabled, device="cpu", init_scale=scale_init, growth_interval=growth_interval) # type: ignore
+        super().__init__(enabled=enabled, device="cpu", init_scale=scale_init, growth_interval=growth_interval)  # type: ignore
         self._enabled = enabled
         self.scale_min = scale_min
         self.device = device
@@ -81,9 +68,7 @@ class GradScaler(torch.amp.GradScaler):
             # We write scale=1 to log if the scaler is disabled
             self._scale = torch.tensor((1,), dtype=torch.float32, device=self.device)
 
-
     def update(self):
-
         if not self._enabled:
             return
 
@@ -105,4 +90,4 @@ def move_optimizer_to(optim, device):
                 if isinstance(subparam, torch.Tensor):
                     subparam.data = subparam.data.to(device)
                     if subparam._grad is not None:
-                        subparam._grad.data = subparam._grad.data.to(device)
+                        subparam._grad.data = subparam._grad.data.to(device)

autogluon/tabular/models/mitra/_internal/core/get_scheduler.py

@@ -1,14 +1,15 @@
 import torch
-from torch.optim.lr_scheduler import ReduceLROnPlateau, LinearLR
+from torch.optim.lr_scheduler import LinearLR, ReduceLROnPlateau
 from transformers import get_constant_schedule_with_warmup
 from transformers.optimization import get_cosine_with_min_lr_schedule_with_warmup
 
 from ..._internal.config.config_pretrain import ConfigPretrain
 
 
-def get_scheduler(hyperparams: dict, optimizer: torch.optim.Optimizer) -> tuple[torch.optim.lr_scheduler.LambdaLR, ReduceLROnPlateau]:
-
-    warmup_steps = hyperparams['warmup_steps']
+def get_scheduler(
+    hyperparams: dict, optimizer: torch.optim.Optimizer
+) -> tuple[torch.optim.lr_scheduler.LambdaLR, ReduceLROnPlateau]:
+    warmup_steps = hyperparams["warmup_steps"]
 
     # if warmup_steps > 0:
     #     scheduler_warmup = torch.optim.lr_scheduler.LambdaLR(
@@ -18,7 +19,7 @@ def get_scheduler(hyperparams: dict, optimizer: torch.optim.Optimizer) -> tuple[
     #     scheduler_warmup = torch.optim.lr_scheduler.LambdaLR(
     #         optimizer, lambda step: 1.0
     #     )
-    
+
     if warmup_steps > 0:
         scheduler_warmup = LinearLR(
             optimizer,
@@ -29,39 +30,23 @@ def get_scheduler(hyperparams: dict, optimizer: torch.optim.Optimizer) -> tuple[
     else:
         scheduler_warmup = torch.optim.lr_scheduler.ConstantLR(optimizer, factor=1.0, total_iters=1)
 
-    if hyperparams['lr_scheduler']:      
+    if hyperparams["lr_scheduler"]:
         scheduler_reduce_on_plateau = ReduceLROnPlateau(
-            optimizer, 
-            patience=hyperparams['lr_scheduler_patience'], 
-            min_lr=0, 
-            factor=0.2
+            optimizer, patience=hyperparams["lr_scheduler_patience"], min_lr=0, factor=0.2
         )
     else:
         # With ReduceLROnPlateau, the scheduler accepts a metric to monitor, so our dummy metric must also be a ReduceLRonPlateau scheduler
-        scheduler_reduce_on_plateau = ReduceLROnPlateau(
-            optimizer, 
-            patience=1000000000, 
-            min_lr=0, 
-            factor=0.2
-        )
+        scheduler_reduce_on_plateau = ReduceLROnPlateau(optimizer, patience=1000000000, min_lr=0, factor=0.2)
 
     return scheduler_warmup, scheduler_reduce_on_plateau
 
 
 def get_scheduler_pretrain(cfg: ConfigPretrain, optimizer: torch.optim.Optimizer):
-
-    
     if cfg.optim.cosine_scheduler:
         schedule = get_cosine_with_min_lr_schedule_with_warmup(
-            optimizer,
-            num_warmup_steps=cfg.optim.warmup_steps,
-            num_training_steps=cfg.optim.steps,
-            min_lr_rate=0.1
+            optimizer, num_warmup_steps=cfg.optim.warmup_steps, num_training_steps=cfg.optim.steps, min_lr_rate=0.1
         )
     else:
-        schedule = get_constant_schedule_with_warmup(
-            optimizer,
-            num_warmup_steps=cfg.optim.warmup_steps
-        )
+        schedule = get_constant_schedule_with_warmup(optimizer, num_warmup_steps=cfg.optim.warmup_steps)
 
-    return schedule
+    return schedule