dynamic imports for log and struc

Author: aditya0by0

chebai/loggers/custom.py

@@ -2,7 +2,6 @@
 from datetime import datetime
 from typing import List, Literal, Optional, Union
 
-import wandb
 from lightning.fabric.utilities.types import _PATH
 from lightning.pytorch.callbacks import ModelCheckpoint
 from lightning.pytorch.loggers import WandbLogger
@@ -105,6 +104,8 @@ def set_fold(self, fold: int) -> None:
         Args:
             fold (int): Cross-validation fold number.
         """
+        import wandb
+
         if fold != self._fold:
             self._fold = fold
             # Start new experiment

chebai/preprocessing/structures.py

@@ -1,8 +1,10 @@
-from typing import Any, Tuple, Union
+from typing import TYPE_CHECKING, Any, Tuple, Union
 
-import networkx as nx
 import torch
 
+if TYPE_CHECKING:
+    import networkx as nx
+
 
 class XYData(torch.utils.data.Dataset):
     """
@@ -129,6 +131,8 @@ def to_x(self, device: torch.device) -> Tuple[nx.Graph, ...]:
         Returns:
             A tuple of molecular graphs with node attributes on the specified device.
         """
+        import networkx as nx
+
         l_ = []
         for g in self.x:
             graph = g.copy()