[src & tests] Fix bug when scheduler=None in System (#223)

asteroid/engine/system.py

@@ -122,10 +122,13 @@ def optimizer_step(
         optimizer_idx: int,
         second_order_closure: Optional[Callable] = None,
     ) -> None:
-        for sched in self.scheduler:
-            if isinstance(sched, dict) and sched["interval"] == "batch":
-                sched["scheduler"].step()  # call step on each batch scheduler
-        self.optimizer.step()
+        if self.scheduler is not None:
+            if not isinstance(self.scheduler, (list, tuple)):
+                self.scheduler = [self.scheduler]  # support multiple schedulers
+            for sched in self.scheduler:
+                if isinstance(sched, dict) and sched["interval"] == "batch":
+                    sched["scheduler"].step()  # call step on each batch scheduler
+            self.optimizer.step()
 
     def validation_step(self, batch, batch_nb):
         """ Need to overwrite PL validation_step to do validation.

tests/engine/system_test.py

@@ -22,3 +22,21 @@ def test_system():
     )
     trainer = Trainer(max_epochs=1, fast_dev_run=True)
     trainer.fit(system)
+
+
+def test_system_no_scheduler():
+    model = nn.Sequential(nn.Linear(10, 10), nn.ReLU())
+    optimizer = optim.Adam(model.parameters(), lr=1e-3)
+    scheduler = None
+    dataset = DummyDataset()
+    loader = data.DataLoader(dataset, batch_size=2, num_workers=4)
+    system = System(
+        model,
+        optimizer,
+        loss_func=nn.MSELoss(),
+        train_loader=loader,
+        val_loader=loader,
+        scheduler=scheduler,
+    )
+    trainer = Trainer(max_epochs=1, fast_dev_run=True)
+    trainer.fit(system)