Skip to content

Commit 328d13b

Browse files
committed
Training of transformers for structure prediction
havily copied from transformers run_squad ..
1 parent 6faf6e3 commit 328d13b

File tree

2 files changed

+722
-2
lines changed

2 files changed

+722
-2
lines changed

sota_extractor2/helpers/training.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,14 @@
11

2-
def set_seed(seed, name, quiet=False):
2+
def set_seed(seed, name, quiet=False, all_gpus=True):
33
import torch
44
import numpy as np
5+
import random
56
if not quiet:
67
print(f"Setting {name} seed to {seed}")
78
torch.manual_seed(seed)
89
torch.backends.cudnn.deterministic = True
910
torch.backends.cudnn.benchmark = False
10-
np.random.seed(seed)
11+
np.random.seed(seed)
12+
random.seed(seed)
13+
if all_gpus:
14+
torch.cuda.manual_seed_all(seed)

0 commit comments

Comments
 (0)