fix: missing shared tensor in bert::BertLMPredictionHead

[silver-ymz]

(cls.predictions.decoder.weight, bert.embeddings.word_embeddings.weight) and (cls.predictions.decoder.bias, cls.predictions.bias) are shared tensor in BertForMaskedLM model. They may not exist in model.safetensors file.

[alvarobartt]

Hi, not sure and better to double check with a maintainer but maybe this syntax is preferred to avoid the contains_tensor call (?)

Suggested change

	let decoder_weight = if vb.contains_tensor("decoder.weight") {
	vb.get((config.vocab_size, config.hidden_size), "decoder.weight")?
	} else {
	vb.root().get(
	(config.vocab_size, config.hidden_size),
	"bert.embeddings.word_embeddings.weight",
	)?
	};
	let decoder_bias = if vb.contains_tensor("decoder.bias") {
	vb.get(config.vocab_size, "decoder.bias")?
	} else {
	vb.get(config.vocab_size, "bias")?
	};
	let decoder_weight = vb
	.get((config.vocab_size, config.hidden_size), "decoder.weight")
	.or_else(|_| {
	vb.root().get(
	(config.vocab_size, config.hidden_size),
	"bert.embeddings.word_embeddings.weight",
	)
	})?;
	let decoder_bias = vb
	.get(config.vocab_size, "decoder.bias")
	.or_else(|_| vb.get(config.vocab_size, "bias"))?;