Some minor changes.

2 files changed, 8 insertions, 4 deletions

diff --git a/src/training/run_experiment.py b/src/training/run_experiment.py
index 55a9572..a883b45 100644
--- a/src/training/run_experiment.py
+++ b/src/training/run_experiment.py
@@ -21,8 +21,9 @@ from training.trainer.train import Trainer
 import wandb
 import yaml
 
-
+import text_recognizer.models
 from text_recognizer.models import Model
+import text_recognizer.networks
 from text_recognizer.networks.loss import loss as custom_loss_module
 
 EXPERIMENTS_DIRNAME = Path(__file__).parents[0].resolve() / "experiments"
@@ -77,13 +78,12 @@ def _load_modules_and_arguments(experiment_config: Dict,) -> Tuple[Callable, Dic
     dataset_ = dataset_args["type"]
 
     # Import the model module and model arguments.
-    models_module = importlib.import_module("text_recognizer.models")
-    model_class_ = getattr(models_module, experiment_config["model"])
+    model_class_ = getattr(text_recognizer.models, experiment_config["model"])
 
     # Import metrics.
     metric_fns_ = (
         {
-            metric: getattr(models_module, metric)
+            metric: getattr(text_recognizer.networks, metric)
             for metric in experiment_config["metrics"]
         }
         if experiment_config["metrics"] is not None
diff --git a/src/training/trainer/train.py b/src/training/trainer/train.py
index 223d9c6..8ae994a 100644
--- a/src/training/trainer/train.py
+++ b/src/training/trainer/train.py
@@ -3,6 +3,7 @@
 from pathlib import Path
 import time
 from typing import Dict, List, Optional, Tuple, Type
+import warnings
 
 from einops import rearrange
 from loguru import logger
@@ -23,6 +24,9 @@ torch.manual_seed(4711)
 torch.cuda.manual_seed(4711)
 
 
+warnings.filterwarnings("ignore")
+
+
 class Trainer:
     """Trainer for training PyTorch models."""