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()