LR schedulers (#573)

allenai · Jul 12, 2023 · 01077eb · 01077eb
1 parent 3a19688
commit 01077eb
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 4 deletions.
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
@@ -20,7 +20,7 @@ env:
   WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
   BEAKER_TOKEN: ${{ secrets.BEAKER_TOKEN }}
   BEAKER_WORKSPACE: ai2/tango-testing
-  BEAKER_DEFAULT_CLUSTER: ai2/tango-gpu-tests
+  BEAKER_DEFAULT_CLUSTER: ai2/allennlp-cirrascale
   BEAKER_IMAGE: petew/tango-testing
   GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
 

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 
 - Removed unnecessary code coverage dev requirements.
+- Fixed issue where new version of torch caused no LR schedulers to be registered.
 
 ## [v1.2.1](https://github.com/allenai/tango/releases/tag/v1.2.1) - 2023-04-06
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -28,7 +28,7 @@ dependencies = [
   "base58",
   "xxhash",
   "filelock>=3.4,<4.0",
-  "click>=8.0,<9.0",
+  "click>=8.0,<8.1.4",
   "click-help-colors>=0.9.1,<0.10",
   "rich>=12.3,<14.0",
   "tqdm>=4.62,<5.0",

diff --git a/tango/integrations/torch/optim.py b/tango/integrations/torch/optim.py
@@ -77,7 +77,7 @@ class LRScheduler(torch.optim.lr_scheduler._LRScheduler, Registrable):
 for name, cls in torch.optim.lr_scheduler.__dict__.items():
     if (
         isinstance(cls, type)
-        and issubclass(cls, torch.optim.lr_scheduler._LRScheduler)
-        and not cls == torch.optim.lr_scheduler._LRScheduler
+        and issubclass(cls, torch.optim.lr_scheduler.LRScheduler)
+        and not cls == torch.optim.lr_scheduler.LRScheduler
     ):
         LRScheduler.register("torch::" + name)(cls)
diff --git a/tango/integrations/transformers/__init__.py b/tango/integrations/transformers/__init__.py
@@ -72,6 +72,8 @@
             transformers::AutoModelForImageSegmentation::from_pretrained
             transformers::AutoModelForInstanceSegmentation::from_config
             transformers::AutoModelForInstanceSegmentation::from_pretrained
+            transformers::AutoModelForMaskGeneration::from_config
+            transformers::AutoModelForMaskGeneration::from_pretrained
             transformers::AutoModelForMaskedImageModeling::from_config
             transformers::AutoModelForMaskedImageModeling::from_pretrained
             transformers::AutoModelForMaskedLM::from_config
@@ -159,6 +161,7 @@
             transformers::inverse_sqrt
             transformers::linear
             transformers::polynomial
+            transformers::reduce_lr_on_plateau
 
 - :class:`~tango.integrations.torch.DataCollator`: All data collators from transformers
   are registered according to their class name (e.g. "transformers::DefaultDataCollator").