From 09f1f458e8a5e381d2357a7391bb6eb51d7122a5 Mon Sep 17 00:00:00 2001
From: Luca Antiga <luca@lightning.ai>
Date: Sun, 8 Jan 2023 15:51:50 +0100
Subject: [PATCH] Move conditional import

---
 train.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/train.py b/train.py
index 2e9379f..5db73a8 100644
--- a/train.py
+++ b/train.py
@@ -74,9 +74,6 @@ if gpu_id == 0:
 torch.manual_seed(1337 + gpu_id) # note: each worker gets a different seed
 torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul
 torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn
-# import wandb conditionally
-if wandb_log:
-    import wandb
 
 # poor man's data loader, TODO evaluate need for actual DataLoader
 data_dir = os.path.join('data', dataset)
@@ -182,6 +179,7 @@ def get_lr(iter):
 
 # logging
 if wandb_log and gpu_id == 0:
+    import wandb
     wandb.init(project=wandb_project, name=wandb_run_name)
     wandb.config = {
         "batch_size": batch_size,