Expected type Float got Half in Loss function

This is my loss model forward:

    def forward(self, sentence_features: Iterable[Dict[str, Tensor]], labels: Tensor):
        rep_a, rep_b = self.model(sentence_features)

        output = cosine_similarity(rep_a.float(), rep_b.float()) # errors on this line unless I add the .float() cast
        loss_fct = nn.MSELoss()
        if labels is not None:
            loss = loss_fct(output, labels.view(-1))
            return loss
        else:
            return [rep_a, rep_b], output

Initialized with:

from apex import amp
model, optimizer = amp.initialize(loss_model, optimizer, opt_level='O1')

Start with the provided forward method, especially the cosine_similarity call, and the amp.initialize setup using opt_level='O1'. Reproduce the Float-versus-Half error and inspect how the loss inputs are typed under mixed precision. Done means the loss model handles the shown inputs without requiring the manual .float() cast.