fix import issue

aramis-lab · Sep 23, 2024 · 00752ff · 00752ff
1 parent 264b003
commit 00752ff
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/clinicadl/optim/optimizer/factory.py b/clinicadl/optim/optimizer/factory.py
@@ -1,5 +1,6 @@
-from typing import Any, Dict, Tuple
+from typing import Any, Dict, Iterable, Iterator, List, Tuple
 
+import torch
 import torch.nn as nn
 import torch.optim as optim
 

diff --git a/tests/unittests/optim/optimizer/test_factory.py b/tests/unittests/optim/optimizer/test_factory.py
@@ -102,10 +102,7 @@ def test_get_optimizer(network):
     assert len(optimizer.param_groups) == 3
 
     # special case : no ELSE mentioned
-    config = OptimizerConfig(
-        optimizer="Adagrad",
-        lr_decay={"conv1": 100},
-    )
+    config = create_optimizer_config("Adagrad")(lr_decay={"conv1": 100})
     optimizer, _ = get_optimizer(network, config)
     assert len(optimizer.param_groups) == 2
     assert optimizer.param_groups[0]["lr_decay"] == 100