Merge pull request #72 from narumiruna/fix-mlconfig-torch

fix mlconfig.torch._register_classes
narumiruna · May 18, 2024 · 2e09fb7 · 2e09fb7
2 parents 0f04eef + c8cfa60
commit 2e09fb7
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/mlconfig/torch/__init__.py b/mlconfig/torch/__init__.py
@@ -19,14 +19,12 @@ def _register_classes(module, superclass, prefix=None, sep="."):
             if prefix is not None:
                 name = prefix + sep + name
 
-            mlconfig.conf(attr, name=name)
+            mlconfig.register(attr, name=name)
 
 
 def get_lr_scheduler_class():
     # PyTorch 2.0 renamed LRScheduler to _LRScheduler
-    name = (
-        "LRScheduler" if hasattr(optim.lr_scheduler, "LRScheduler") else "_LRScheduler"
-    )
+    name = "LRScheduler" if hasattr(optim.lr_scheduler, "LRScheduler") else "_LRScheduler"
     return getattr(optim.lr_scheduler, name)
 
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -31,9 +31,9 @@ select = [
     "E", # pycodestyle errors
     "F", # pyflakes
     "I", # isort
+    "SIM", # flake8-simplify
     # "UP", # pyupgrade
     "W", # pycodestyle warnings
-
 ]
 
 [tool.ruff.lint.per-file-ignores]