optimizer and scheduler classes

tests/test_optimizer.py

@@ -0,0 +1,20 @@
+
+import torch
+import pytest
+from pina import TorchOptimizer
+
+opt_list = [
+    torch.optim.Adam,
+    torch.optim.AdamW,
+    torch.optim.SGD,
+    torch.optim.RMSprop
+]
+
+@pytest.mark.parametrize("optimizer_class", opt_list)
+def test_constructor(optimizer_class):
+    TorchOptimizer(optimizer_class, lr=1e-3)
+
+@pytest.mark.parametrize("optimizer_class", opt_list)
+def test_hook(optimizer_class):
+    opt = TorchOptimizer(optimizer_class, lr=1e-3)
+    opt.hook(torch.nn.Linear(10, 10).parameters())

tests/test_scheduler.py

@@ -0,0 +1,27 @@
+
+import torch
+import pytest
+from pina import TorchOptimizer, TorchScheduler
+
+opt_list = [
+    torch.optim.Adam,
+    torch.optim.AdamW,
+    torch.optim.SGD,
+    torch.optim.RMSprop
+]
+
+sch_list = [
+    torch.optim.lr_scheduler.ConstantLR
+]
+
+@pytest.mark.parametrize("scheduler_class", sch_list)
+def test_constructor(scheduler_class):
+    TorchScheduler(scheduler_class)
+
+@pytest.mark.parametrize("optimizer_class", opt_list)
+@pytest.mark.parametrize("scheduler_class", sch_list)
+def test_hook(optimizer_class, scheduler_class):
+    opt = TorchOptimizer(optimizer_class, lr=1e-3)
+    opt.hook(torch.nn.Linear(10, 10).parameters())
+    sch = TorchScheduler(scheduler_class)
+    sch.hook(opt)