You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
11 lines
244 B
11 lines
244 B
1 month ago
|
from tinygrad import Tensor, nn
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
vocab_size = 50257
|
||
|
|
n_embd = 768
|
||
|
|
lm_head = nn.Linear(n_embd, vocab_size, bias=False)
|
||
|
|
bs = 4
|
||
|
|
seq_len = 1024
|
||
|
|
x = Tensor.rand(bs, seq_len, n_embd)
|
||
|
|
ret = lm_head(x).realize()
|