def maybe_convert_weights()

in summarize_from_feedback/models/transformer.py [0:0]

• 11 lines of code
• 5 McCabe index (conditional complexity)

def maybe_convert_weights(model, *, fp16_conv_weights: bool, fp16_embedding_weights: bool):
    """
    Scans through the parameters in the model and converts Conv1D and/or Embedding parameters to
    float16 depending on fp16_conv_weights and fp16_embedding_weights.
    """
    if fp16_conv_weights:

        def _convert_conv_weights_to_fp16(l):
            if isinstance(l, Conv1D):
                l.weight.data = l.weight.data.half()

        model.apply(_convert_conv_weights_to_fp16)

    if fp16_embedding_weights:

        def _convert_embedding_weights_to_fp16(l):
            if isinstance(l, torch.nn.Embedding):
                l.weight.data = l.weight.data.half()

        model.apply(_convert_embedding_weights_to_fp16)