🔀 [Merge] branch 'MODELv2' into INFERENCE

yolo/lazy.py

@@ -26,7 +26,7 @@ def main(cfg: Config):
         model = FastModelLoader(cfg).load_model()
         device = torch.device(cfg.device)
     else:
-        model = create_model(cfg.model, class_num=cfg.class_num, weight_path=cfg.weight).to(device)
+        model = create_model(cfg.model, class_num=cfg.class_num, weight_path=cfg.weight, device=device)
 
     vec2box = Vec2Box(model, cfg.image_size, device)
 

yolo/model/yolo.py

@@ -2,9 +2,9 @@ import os
 from typing import Any, Dict, List, Union
 
 import torch
-import torch.nn as nn
 from loguru import logger
 from omegaconf import ListConfig, OmegaConf
+from torch import device, nn
 
 from yolo.config.config import Config, ModelConfig, YOLOLayer
 from yolo.tools.dataset_preparation import prepare_weight
@@ -117,7 +117,9 @@ class YOLO(nn.Module):
             raise ValueError(f"Unsupported layer type: {layer_type}")
 
 
-def create_model(model_cfg: ModelConfig, class_num: int = 80, weight_path: str = "weights/v9-c.pt") -> YOLO:
+def create_model(
+    model_cfg: ModelConfig, class_num: int = 80, weight_path: str = "weights/v9-c.pt", device: device = device("cuda")
+) -> YOLO:
     """Constructs and returns a model from a Dictionary configuration file.
 
     Args:
@@ -133,9 +135,9 @@ def create_model(model_cfg: ModelConfig, class_num: int = 80, weight_path: str =
         if not os.path.exists(weight_path):
             logger.info(f"🌐 Weight {weight_path} not found, try downloading")
             prepare_weight(weight_path=weight_path)
-        model.model.load_state_dict(torch.load(weight_path))
+        model.model.load_state_dict(torch.load(weight_path, map_location=device))
         logger.info("✅ Success load model weight")
 
     log_model_structure(model.model)
     draw_model(model=model)
-    return model
+    return model.to(device)