Fix: Variable name

[dhdbsrlw]

Hello.

I think variable's name should be changed in 'lit_llama/utils.py.'

When I implement your code, the existing code causes error. Then I debugged with printing.

There was no key named 'transformer.wte.weight' in checkpoint.

Thank you.

[rasbt]

I think the original key should exist for the original Llama weights. Maybe you downloaded the weights from somewhere else? Perhaps the following would be a good compromise:

Suggested change

	embedding_size = checkpoint['lm_head.weight'].shape[1]
	if 'transformer.wte.weight' in checkpoint:
	embedding_size = checkpoint['transformer.wte.weight'].shape[1]
	elif 'lm_head.weight' in checkpoint:
	embedding_size = checkpoint['lm_head.weight'].shape[1]
	else:
	raise ValueError("Neither 'transformer.wte.weight' nor 'lm_head.weight' found in the checkpoint")

[rasbt]

Can't reproduce it at the moment but perhaps we can have a compromise here that remains backwards compatible.