update task

flash/core/schedulers.py

@@ -4,7 +4,7 @@
 _SCHEDULER_REGISTRY = FlashRegistry("scheduler")
 
 if _TRANSFORMERS_AVAILABLE:
-    from transformers.optimization import TYPE_TO_SCHEDULER_FUNCTION
-
-    for v in TYPE_TO_SCHEDULER_FUNCTION.values():
-        _SCHEDULER_REGISTRY(v, name=v.__name__[4:])
+    from transformers import optimization
+    functions = [getattr(optimization, n) for n in dir(optimization) if ("get_" in n and n != 'get_scheduler')]
+    for fn in functions:
+        _SCHEDULER_REGISTRY(fn, name=fn.__name__[4:])

tests/core/test_model.py

@@ -197,7 +197,7 @@ def test_optimization(tmpdir):
     assert isinstance(scheduler[0], torch.optim.lr_scheduler.StepLR)
 
     if _TRANSFORMERS_AVAILABLE:
-        from transformers.optimization import SchedulerType, TYPE_TO_SCHEDULER_FUNCTION
+        from transformers.optimization import get_linear_schedule_with_warmup
 
         assert task.available_schedulers() == [
             'constant_schedule', 'constant_schedule_with_warmup', 'cosine_schedule_with_warmup',
@@ -219,11 +219,11 @@ def test_optimization(tmpdir):
             scheduler_kwargs={"num_warmup_steps": 0.1},
             loss_fn=F.nll_loss,
         )
-        trainer = flash.Trainer(max_epochs=100)
+        trainer = flash.Trainer(max_epochs=1, limit_train_batches=2)
         ds = DummyDataset()
         trainer.fit(task, train_dataloader=DataLoader(ds))
         optimizer, scheduler = task.configure_optimizers()
         assert isinstance(optimizer[0], torch.optim.Adadelta)
         assert isinstance(scheduler[0], torch.optim.lr_scheduler.LambdaLR)
-        expected = TYPE_TO_SCHEDULER_FUNCTION[SchedulerType.LINEAR].__name__
+        expected = get_linear_schedule_with_warmup.__name__
         assert scheduler[0].lr_lambdas[0].__qualname__.split('.')[0] == expected