finally?

models/__init__.py

@@ -21,6 +21,12 @@
 import numpy as np
 import yaml
 
+# Import torch for GPU memory monitoring
+try:
+    import torch
+except ImportError:
+    torch = None
+
 # --------------------------------------------------------------------------------------
 # Logging (ensure a handler exists very early)
 # --------------------------------------------------------------------------------------