osmarks/nanogpt-experiments
1
0
Fork 0
mirror of https://github.com/osmarks/nanogpt-experiments.git synced 2025-09-01 10:27:58 +00:00
Code Issues Releases Wiki Activity

Move conditional import

Browse Source
This commit is contained in:
Luca Antiga
2023-01-08 15:51:50 +01:00
parent aba47f0a35
commit 09f1f458e8
1 changed files with 1 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
train.py
View File

@@ -74,9 +74,6 @@ if gpu_id == 0:
torch.manual_seed(1337 + gpu_id) # note: each worker gets a different seed 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.cuda.matmul.allow_tf32 = True # allow tf32 on matmul
torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn 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 # poor man's data loader, TODO evaluate need for actual DataLoader
data_dir = os.path.join('data', dataset) data_dir = os.path.join('data', dataset)
@@ -182,6 +179,7 @@ def get_lr(iter):
# logging # logging
if wandb_log and gpu_id == 0: if wandb_log and gpu_id == 0:
import wandb
wandb.init(project=wandb_project, name=wandb_run_name) wandb.init(project=wandb_project, name=wandb_run_name)
wandb.config = { wandb.config = {
"batch_size": batch_size, "batch_size": batch_size,