Added lazy loading for torch

Author: m-messer

evaluation_function/lazy_load.py

@@ -0,0 +1,19 @@
+import importlib
+
+
+class LazyModule:
+    """
+    This class is response for lazy loading of heavy imports
+    """
+    def __init__(self, name):
+        self._name = name
+        self._module = None
+
+    def _load(self):
+        if self._module is None:
+            self._module = importlib.import_module(self._name)
+        return self._module
+
+    def __getattr__(self, item):
+        return getattr(self._load(), item)
+

evaluation_function/models/basic_nn.py

@@ -16,21 +16,23 @@
 
 """
 
-import torch
-import torch.nn as nn
-import torch.optim as optim
-
 from lf_toolkit.evaluation import Result, Params
 
 from pathlib import Path
 import os
 
+from evaluation_function.lazy_load import LazyModule
+
 # Setup paths for saving/loading model and data
 BASE_DIR = Path(__file__).resolve().parent
 MODEL_DIR = Path(os.environ.get("MODEL_DIR", BASE_DIR / "storage"))
 MODEL_DIR.mkdir(parents=True, exist_ok=True)
 MODEL_PATH = MODEL_DIR / "basic_nn.pt"
 
+torch = LazyModule("torch")
+nn = LazyModule("torch.nn")
+optim = LazyModule("torch.optim")
+
 def f(x):
     """Target function with noise (sine wave)."""
     return torch.sin(x)