[Question] Why is the call to allocate_workspace() using only the first input as size parameters?

In DeepSpeedTransformerInference, the function self.allocate_workspace() is called only once per instance, on the first forward() pass call.

This call is made using the batch size (input.size()[0]) and the sequence size (input.size()[1]) of the first forward() pass:
https://github.com/microsoft/DeepSpeed/blob/901d80701258f59fc951c273f4c488ca04b10867/deepspeed/model_implementations/transformers/ds_transformer.py#L140-L147

My question is why? Either:

• The DeepSpeedTransformerInference instance is meant to be reused with only inputs up to these sizes, otherwise risking to exceed the workspace memory. If this is the case, I haven't see this documented anywhere
• DeepSpeedTransformerInference instances are meant to be used only once, and in that case I have not seen that documented anywhere either

Start in deepspeed/model_implementations/transformers/ds_transformer.py at the allocate_workspace() call around lines 140-147, and trace how the first forward() inputs determine workspace sizing. Check whether reuse and larger subsequent inputs are documented elsewhere; done means documenting the confirmed instance-reuse and size behavior.