Merge pull request #160 from ArgoHA/master

ruff 100 formatting

reference/convert_weight.py

@@ -1,14 +1,16 @@
-import torch
-import os
 import argparse
+import os
+
+import torch
+
 
 def save_only_ema_weights(checkpoint_file):
     """Extract and save only the EMA weights."""
-    checkpoint = torch.load(checkpoint_file, map_location='cpu')
+    checkpoint = torch.load(checkpoint_file, map_location="cpu")
 
     weights = {}
-    if 'ema' in checkpoint:
-        weights['model'] = checkpoint['ema']['module']
+    if "ema" in checkpoint:
+        weights["model"] = checkpoint["ema"]["module"]
     else:
         raise ValueError("The checkpoint does not contain 'ema'.")
 
@@ -19,9 +21,10 @@ def save_only_ema_weights(checkpoint_file):
     torch.save(weights, output_file)
     print(f"EMA weights saved to {output_file}")
 
-if __name__ == '__main__':
+
+if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Extract and save only EMA weights.")
-    parser.add_argument('checkpoint_file', type=str, help="Path to the input checkpoint file.")
+    parser.add_argument("checkpoint_file", type=str, help="Path to the input checkpoint file.")
 
     args = parser.parse_args()
     save_only_ema_weights(args.checkpoint_file)

src/__init__.py

@@ -3,7 +3,4 @@
 """
 
 # for register purpose
-from . import optim
-from . import data
-from . import nn
-from . import zoo
+from . import data, nn, optim, zoo

src/core/__init__.py

@@ -3,7 +3,7 @@
 Copyright(c) 2023 lyuwenyu. All Rights Reserved.
 """
 
-from .workspace import GLOBAL_CONFIG, register, create
-from .yaml_utils import *
 from ._config import BaseConfig
+from .workspace import GLOBAL_CONFIG, create, register
 from .yaml_config import YAMLConfig
+from .yaml_utils import *