src/transformers/models/gptj/modeling_tf_gptj.py/TFGPTJMainLayer/set_input_embeddings
class TFGPTJMainLayer:    def set_input_embeddings(self, value: tf.Tensor):
        self.wte.weight = value
        self.wte.vocab_size = shape_list(value)[0]