From 6c9ba1d8fc3f0957ea7a1a59ef8622f8d54ac89a Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Tue, 4 Aug 2020 19:22:43 +0200 Subject: [PATCH] [Reformer] Make random seed generator available on random seed and not on model device (#6244) * improve if else statement random seeds * Apply suggestions from code review * Update src/transformers/modeling_reformer.py --- src/transformers/modeling_reformer.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/src/transformers/modeling_reformer.py b/src/transformers/modeling_reformer.py index f9e8ac76d6..6049dc3ed0 100644 --- a/src/transformers/modeling_reformer.py +++ b/src/transformers/modeling_reformer.py @@ -1399,15 +1399,16 @@ class ReformerLayer(nn.Module): """ # randomize seeds - if next(self.parameters()).device.type == "cuda": + # use cuda generator if available + if len(torch.cuda.default_generators) > 0: # GPU device_idx = torch.cuda.current_device() self.attention_seed = torch.cuda.default_generators[device_idx].seed() - torch.cuda.manual_seed(self.attention_seed) else: # CPU self.attention_seed = int(torch.seed() % sys.maxsize) - torch.manual_seed(self.attention_seed) + + torch.manual_seed(self.attention_seed) def _init_feed_forward_seed(self): """ @@ -1417,17 +1418,17 @@ class ReformerLayer(nn.Module): call and 1 forward call in backward to recalculate activations. """ - # randomize seeds - if next(self.parameters()).device.type == "cuda": + # use cuda generator if available + if len(torch.cuda.default_generators) > 0: # GPU device_idx = torch.cuda.current_device() self.feed_forward_seed = torch.cuda.default_generators[device_idx].seed() - torch.cuda.manual_seed(self.feed_forward_seed) else: # CPU self.feed_forward_seed = int(torch.seed() % sys.maxsize) - torch.manual_seed(self.feed_forward_seed) + + torch.manual_seed(self.feed_forward_seed) def forward( self,