Skip to content

Commit 1d90341

Browse files
malfetSvetlana Karslioglu
and
Svetlana Karslioglu
authored
[BE] Simplify ids_tensor (#2431)
Remove `global_rng` and use `torch.randint` to feel the tensor of shape `shape` with values in range `[0, vocab_size)` Co-authored-by: Svetlana Karslioglu <svekars@fb.com>
1 parent d541f74 commit 1d90341

File tree

2 files changed

+4
-30
lines changed

2 files changed

+4
-30
lines changed

intermediate_source/dynamic_quantization_bert_tutorial.rst

Lines changed: 2 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -253,8 +253,6 @@ model before and after the dynamic quantization.
253253
torch.manual_seed(seed)
254254
set_seed(42)
255255
256-
# Initialize a global random number generator
257-
global_rng = random.Random()
258256
259257
260258
2.2 Load the fine-tuned BERT model
@@ -526,20 +524,9 @@ We can serialize and save the quantized model for the future use using
526524

527525
.. code:: python
528526
529-
def ids_tensor(shape, vocab_size, rng=None, name=None):
527+
def ids_tensor(shape, vocab_size):
530528
# Creates a random int32 tensor of the shape within the vocab size
531-
if rng is None:
532-
rng = global_rng
533-
534-
total_dims = 1
535-
for dim in shape:
536-
total_dims *= dim
537-
538-
values = []
539-
for _ in range(total_dims):
540-
values.append(rng.randint(0, vocab_size - 1))
541-
542-
return torch.tensor(data=values, dtype=torch.long, device='cpu').view(shape).contiguous()
529+
return torch.randint(0, vocab_size, shape=shape, dtype=torch.int, device='cpu')
543530
544531
input_ids = ids_tensor([8, 128], 2)
545532
token_type_ids = ids_tensor([8, 128], 2)

prototype_source/graph_mode_dynamic_bert_tutorial.rst

Lines changed: 2 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -60,22 +60,9 @@ Once all the necesessary packages are downloaded and installed we setup the code
6060
from torch.quantization import per_channel_dynamic_qconfig
6161
from torch.quantization import quantize_dynamic_jit
6262
63-
global_rng = random.Random()
64-
65-
def ids_tensor(shape, vocab_size, rng=None, name=None):
63+
def ids_tensor(shape, vocab_size):
6664
# Creates a random int32 tensor of the shape within the vocab size
67-
if rng is None:
68-
rng = global_rng
69-
70-
total_dims = 1
71-
for dim in shape:
72-
total_dims *= dim
73-
74-
values = []
75-
for _ in range(total_dims):
76-
values.append(rng.randint(0, vocab_size - 1))
77-
78-
return torch.tensor(data=values, dtype=torch.long, device='cpu').view(shape).contiguous()
65+
return torch.randint(0, vocab_size, shape=shape, dtype=torch.int, device='cpu')
7966
8067
# Setup logging
8168
logger = logging.getLogger(__name__)

0 commit comments

Comments
 (0)