[Proposal] Make wandb a lazy import to avoid unnecessary overhead

[liusida]

Proposal

TransformerLens/transformer_lens/train.py

Line 12 in 50ee38b

import wandb

Move import wandb from top-level in train.py (so it always executes on import) into inside the train() function (and guard with if cfg.wandb)
Reason: wandb is large and slow to import; many consumers of TransformerLens don’t train and thus don’t need wandb, so avoiding the import reduces overhead on slower machines.

[liusida]

Related to this:

wandb/wandb#10898
[Bug]: import wandb may take ~100 seconds when sys.modules contains many entries