#256 Add netspresso_trainer lazy import (#257)

Nota-NetsPresso · Jun 24, 2024 · a58a126 · a58a126
1 parent f051995
commit a58a126
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/netspresso/trainer/trainer.py b/netspresso/trainer/trainer.py
@@ -2,7 +2,6 @@
 from typing import Any, Dict, List, Optional, Union
 
 from loguru import logger
-from netspresso_trainer import train_with_yaml
 from omegaconf import OmegaConf
 
 from netspresso.enums import Status, Task, TaskType
@@ -361,6 +360,8 @@ def train(self, gpus: str, project_name: str) -> Dict:
             Dict: A dictionary containing information about the training.
         """
 
+        from netspresso_trainer import train_with_yaml
+
         self._validate_config()
         self._apply_img_size()