Error loading model via from_pretrained

I have trained a TFDistilBertForSequenceClassification model and successfully saved it to disk using save_pretrained. The expected files (tf_model.h5 and config.json) are in my model directory. However, when I try to load this model via TFDistilBertForSequenceClassification.from_pretrained, I get the following error concerning embeddings.


InvalidArgumentError Traceback (most recent call last)
in
1 model = TFDistilBertForSequenceClassification.from_pretrained(
2 ckpt_dir,
----> 3 config=distilbertconfig
4 )

/opt/anaconda3/envs/subreg/lib/python3.7/site-packages/transformers/modeling_tf_utils.py in from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs)
1323 model(model.dummy_inputs) # build the network with dummy inputs
1324 else:
→ 1325 model(model.dummy_inputs) # build the network with dummy inputs
1326
1327 assert os.path.isfile(resolved_archive_file), f"Error retrieving file {resolved_archive_file}"

/opt/anaconda3/envs/subreg/lib/python3.7/site-packages/keras/utils/traceback_utils.py in error_handler(*args, **kwargs)
65 except Exception as e: # pylint: disable=broad-except
66 filtered_tb = _process_traceback_frames(e.traceback)
—> 67 raise e.with_traceback(filtered_tb) from None
68 finally:
69 del filtered_tb

/opt/anaconda3/envs/subreg/lib/python3.7/site-packages/transformers/models/distilbert/modeling_tf_distilbert.py in call(self, input_ids, attention_mask, head_mask, inputs_embeds, output_attentions, output_hidden_states, return_dict, labels, training, **kwargs)
806 output_hidden_states=inputs[“output_hidden_states”],
807 return_dict=inputs[“return_dict”],
→ 808 training=inputs[“training”],
809 )
810 hidden_state = distilbert_output[0] # (bs, seq_len, dim)

/opt/anaconda3/envs/subreg/lib/python3.7/site-packages/transformers/models/distilbert/modeling_tf_distilbert.py in call(self, input_ids, attention_mask, head_mask, inputs_embeds, output_attentions, output_hidden_states, return_dict, training, **kwargs)
414
415 embedding_output = self.embeddings(
→ 416 inputs[“input_ids”], inputs_embeds=inputs[“inputs_embeds”]
417 ) # (bs, seq_length, dim)
418 tfmr_output = self.transformer(

/opt/anaconda3/envs/subreg/lib/python3.7/site-packages/transformers/models/distilbert/modeling_tf_distilbert.py in call(self, input_ids, position_ids, inputs_embeds, training)
110
111 if input_ids is not None:
→ 112 inputs_embeds = tf.gather(params=self.weight, indices=input_ids)
113
114 input_shape = shape_list(inputs_embeds)[:-1]

InvalidArgumentError: Exception encountered when calling layer “embeddings” (type TFEmbeddings).

indices[0,0] = 7 is not in [0, 5) [Op:ResourceGather]

Call arguments received:
• input_ids=tf.Tensor(shape=(3, 5), dtype=int32)
• position_ids=None
• inputs_embeds=None
• training=False

Potentially noteworthy is that I got the following warning when saving my model using save_pretrained:

WARNING:absl:Found untraced functions such as embeddings_layer_call_fn, embeddings_layer_call_and_return_conditional_losses, transformer_layer_call_fn, transformer_layer_call_and_return_conditional_losses, add_layer_call_fn while saving (showing 5 of 62). These functions will not be directly callable after loading.

I read online that this warning is harmless, but the warning mentions embeddings, so I figured it may be related to my error.

I would very much appreciate any references or guidance on how to properly load my TFDistilBertForSequenceClassification model after it is saved. Thank you!