fix_batch_precision_for_deepspeed

composer.distributed.fix_batch_precision_for_deepspeed(batch, precision)

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

Note

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.

Parameters

• batch (Batch) – The batch of data to adjust the precision for.

• precision (Precision) – The precision to use.

Returns

Batch – The batch with it’s precision adjusted to the specified precision.