Fix LRScheduler issue of PyTorch (#241)

* Fix LRScheduler issue of PyTorch

- PyTorch 2.0.0 has deprecated _LRScheduler, use LRScheduler instead
- Imported _LRScheduler or LRScheduler as PyTorchLRScheduler

* Fix lint

Author: theory-in-progress

src/templates/template-text-classification/main.py

@@ -11,10 +11,14 @@
 from ignite.utils import manual_seed
 from models import TransformerModel
 from torch import nn, optim
-from torch.optim.lr_scheduler import _LRScheduler
 from trainers import setup_evaluator, setup_trainer
 from utils import *
 
+try:
+    from torch.optim.lr_scheduler import _LRScheduler as PyTorchLRScheduler
+except ImportError:
+    from torch.optim.lr_scheduler import LRScheduler as PyTorchLRScheduler
+
 os.environ["TOKENIZERS_PARALLELISM"] = "false"  # remove tokenizer paralleism warning
 
 
@@ -78,10 +82,10 @@ def run(local_rank: int, config: Any):
     (config.output_dir / "config-lock.yaml").write_text(yaml.dump(config))
     trainer.logger = evaluator.logger = logger
 
-    if isinstance(lr_scheduler, _LRScheduler):
+    if isinstance(lr_scheduler, PyTorchLRScheduler):
         trainer.add_event_handler(
             Events.ITERATION_COMPLETED,
-            lambda engine: cast(_LRScheduler, lr_scheduler).step(),
+            lambda engine: cast(PyTorchLRScheduler, lr_scheduler).step(),
         )
     elif isinstance(lr_scheduler, LRScheduler):
         trainer.add_event_handler(Events.ITERATION_COMPLETED, lr_scheduler)

src/templates/template-vision-segmentation/main.py

@@ -11,11 +11,16 @@
 from ignite.utils import manual_seed
 from models import setup_model
 from torch import nn, optim
-from torch.optim.lr_scheduler import _LRScheduler, LambdaLR
+from torch.optim.lr_scheduler import LambdaLR
 from trainers import setup_evaluator, setup_trainer
 from utils import *
 from vis import predictions_gt_images_handler
 
+try:
+    from torch.optim.lr_scheduler import LRScheduler as PyTorchLRScheduler
+except ImportError:
+    from torch.optim.lr_scheduler import _LRScheduler as PyTorchLRScheduler
+
 
 def run(local_rank: int, config: Any):
     # make a certain seed
@@ -71,10 +76,10 @@ def run(local_rank: int, config: Any):
     (config.output_dir / "config-lock.yaml").write_text(yaml.dump(config))
     trainer.logger = evaluator.logger = logger
 
-    if isinstance(lr_scheduler, _LRScheduler):
+    if isinstance(lr_scheduler, PyTorchLRScheduler):
         trainer.add_event_handler(
             Events.ITERATION_COMPLETED,
-            lambda engine: cast(_LRScheduler, lr_scheduler).step(),
+            lambda engine: cast(PyTorchLRScheduler, lr_scheduler).step(),
         )
     elif isinstance(lr_scheduler, LRScheduler):
         trainer.add_event_handler(Events.ITERATION_COMPLETED, lr_scheduler)