Learning rate scheduler can be checkpointed (#218, Closes #217)

HasStateDict type changed to include torch.optim.lr_scheduler._LRScheduler which was missing before, causing the checkpointer to not save/load the state_dict of LRSchedulers.
NKI-AI · Jul 13, 2022 · c058f4c · c058f4c
1 parent 702b8f4
commit c058f4c
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/direct/nn/openvino/openvino_model.py b/direct/nn/openvino/openvino_model.py
@@ -1,9 +1,9 @@
 import io
 
 import torch
-from torch import nn
 from openvino.inference_engine import IECore
 from openvino_extensions import get_extensions_path
+from torch import nn
 
 
 class InstanceNorm2dFunc(torch.autograd.Function):

diff --git a/direct/types.py b/direct/types.py
@@ -10,4 +10,4 @@
 Number = Union[float, int]
 PathOrString = Union[pathlib.Path, str]
 FileOrUrl = NewType("FileOrUrl", PathOrString)
-HasStateDict = Union[nn.Module, torch.optim.Optimizer, GradScaler]
+HasStateDict = Union[nn.Module, torch.optim.Optimizer, torch.optim.lr_scheduler._LRScheduler, GradScaler]
diff --git a/setup.py b/setup.py
@@ -39,7 +39,7 @@
         "numpy>=1.21.2",
         "h5py==3.3.0",
         "omegaconf==2.1.1",
-        "torch>=1.10.2",
+        "torch==1.11.0",
         "torchvision",
         "scikit-image>=0.19.0",
         "scikit-learn>=1.0.1",