Replace python random with torch.rand to enable dynamo.export (#24434)

* Replace python random with torch.rand to enable dynamo.export

* revert changes to flax model code

* Remove unused random import

* Fix torch template

* Move torch.manual_seed(0) to right location
This commit is contained in:
Bowen Bao
2023-06-23 05:17:21 -07:00
committed by GitHub
parent c036c814f4
commit a28325e25e
42 changed files with 71 additions and 104 deletions

View File

@@ -1560,7 +1560,6 @@ class {{cookiecutter.camelcase_modelname}}ForQuestionAnswering({{cookiecutter.ca
{% else %}
import math
import copy
import random
from typing import Optional, Tuple, List, Union
import torch
@@ -2306,7 +2305,7 @@ class {{cookiecutter.camelcase_modelname}}Encoder({{cookiecutter.camelcase_model
if output_hidden_states:
encoder_states = encoder_states + (hidden_states,)
# add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
dropout_probability = random.uniform(0, 1)
dropout_probability = torch.randn([])
if self.training and (dropout_probability < self.layerdrop): # skip the layer
layer_outputs = (None, None)
else:
@@ -2543,7 +2542,7 @@ class {{cookiecutter.camelcase_modelname}}Decoder({{cookiecutter.camelcase_model
# add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
if output_hidden_states:
all_hidden_states += (hidden_states,)
dropout_probability = random.uniform(0, 1)
dropout_probability = torch.randn([])
if self.training and (dropout_probability < self.layerdrop):
continue