composer.distributed.fix_batch_precision_for_deepspeed(batch, precision)[source]#

Ensures that a batch is properly formatted for DeepSpeed precisions, if active.


Just because the precision is set to FP16 doesnโ€™t mean the entire batch can be FP16 too. For example, integer tensors are common in inputs and outputs of various models, and these must not be converted. The assumption here is that a tensor should only be converted to FP16 if it was given in FP32.

  • batch (Batch) โ€“ The batch of data to adjust the precision for.

  • precision (Precision) โ€“ The precision to use.


Batch โ€“ The batch with itโ€™s precision adjusted to the specified precision.