slight import refactor for rank_zero

speediedan · Aug 15, 2024 · 2e7be86 · 2e7be86
1 parent 2b35455
commit 2e7be86
Show file tree

Hide file tree

Showing 5 changed files with 5 additions and 6 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,7 +4,7 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
-## [2.4.0] - 2024-XX-XX
+## [2.4.0] - 2024-08-15
 
 ### Added
 

diff --git a/requirements/base.txt b/requirements/base.txt
@@ -2,4 +2,3 @@
 # the below is uncommented when master is targeting a specific pl dev master commit
 git+https://github.com/Lightning-AI/lightning.git@2064887b12dd934a5f9a2bf45897f29e3bfc74d1#egg=lightning
 torch>=2.1.0
-#mpmath<1.4.0  # temporary requirement to avoid installation of alpha version of mpmath
diff --git a/requirements/standalone_base.txt b/requirements/standalone_base.txt
@@ -2,4 +2,3 @@
 # the below is uncommented when master is targeting a specific pl dev master commit
 git+https://github.com/Lightning-AI/pytorch-lightning.git@2064887b12dd934a5f9a2bf45897f29e3bfc74d1#egg=pytorch-lightning
 torch>=2.1.0
-#mpmath<1.4.0  # temporary requirement to avoid installation of alpha version of mpmath
diff --git a/src/finetuning_scheduler/fts.py b/src/finetuning_scheduler/fts.py
@@ -24,13 +24,13 @@
 
 import lightning.pytorch as pl
 import torch
-from lightning.fabric.utilities import rank_zero_info
+from lightning.fabric.utilities import rank_zero_info, rank_zero_warn
 from lightning.fabric.utilities.distributed import ReduceOp
 from lightning.pytorch.callbacks import BaseFinetuning
 from lightning.pytorch.strategies.strategy import Strategy
 from lightning.pytorch.trainer.states import TrainerFn
 from lightning.pytorch.utilities.exceptions import MisconfigurationException
-from lightning.pytorch.utilities.rank_zero import rank_zero_debug, rank_zero_warn
+from lightning.pytorch.utilities.rank_zero import rank_zero_debug
 
 from finetuning_scheduler.fts_supporters import (
     CallbackDepMixin,

diff --git a/src/finetuning_scheduler/strategy_adapters/base.py b/src/finetuning_scheduler/strategy_adapters/base.py
@@ -22,12 +22,13 @@
 
 import torch
 
-from lightning.fabric.utilities import rank_zero_info, rank_zero_debug
+from lightning.fabric.utilities import rank_zero_info
 from lightning.fabric.utilities.types import ReduceLROnPlateau
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks import Callback
 from lightning.pytorch.callbacks import BaseFinetuning
 from lightning.pytorch.strategies.strategy import Strategy
+from lightning.pytorch.utilities.rank_zero import rank_zero_debug
 
 
 class StrategyAdapter: