github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[NFC] polish tests/test_layers/test_sequence/checks_seq/check_layer_seq.py code style (#1723)

pull/1743/head
xyupeng 2022-10-18 14:16:40 +08:00 committed by Frank Lee
parent 7e62af28a0
commit ff373a11eb
1 changed files with 3 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
tests/test_layers/test_sequence/checks_seq/check_layer_seq.py
View File

@ -12,15 +12,10 @@ def check_selfattention():
BATCH = 4 BATCH = 4
HIDDEN_SIZE = 16 HIDDEN_SIZE = 16
layer = TransformerSelfAttentionRing( layer = TransformerSelfAttentionRing(16, 8, 8, 0.1)
16,
8,
8,
0.1
)
layer = layer.to(get_current_device()) layer = layer.to(get_current_device())
hidden_states = torch.rand(SUB_SEQ_LENGTH, BATCH, HIDDEN_SIZE).to(get_current_device()) hidden_states = torch.rand(SUB_SEQ_LENGTH, BATCH, HIDDEN_SIZE).to(get_current_device())
attention_mask = torch.randint(low=0, high=2, size=(BATCH, 1, 1, 1, SUB_SEQ_LENGTH * WORLD_SIZE)).to( attention_mask = torch.randint(low=0, high=2,
get_current_device()) size=(BATCH, 1, 1, 1, SUB_SEQ_LENGTH * WORLD_SIZE)).to(get_current_device())
out = layer(hidden_states, attention_mask) out = layer(hidden_states, attention_mask)