Thanks for the reply. This method is good and useful for a sequence input, but I was more looking for something that could resize the matrices that depend on the vocab_size of a transformer.
e.g. the embedding matrix of a transformer usually takes on dimensions something like (vocab_size, 1024), where the vocab_size might be something like 52153. A matrix of this size isn’t efficient to pass onto the tensor cores, so I was looking for a way to pad it so it was a multiple of 8 (e.g. to 52160).