Hi @ptrblck! Back again with another question… 
Not sure if I should’ve opened a new thread, so here we go:
I have this piece of code (a attention calculation) that I am trying to fix while using mixed precision (for the model):
    def att(self, value, key, query, mask):
        d_k = query.size(-1)

        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)

        if mask is not None:
            # print(scores.shape, mask.shape, torch.finfo(scores.dtype).min)
            scores = scores.masked_fill(mask, torch.finfo(scores.dtype).min)

        att_map = F.softmax(scores, dim=-1)
        att_map = self.dropout(att_map)

        return torch.matmul(att_map, value)

This now works for me when using non-mixed-precision, however it does not when using mixed-precision. This had to do with the softmax function (and the mask not being able to use sufficiently low numbers with fp16, I think) that gives a overflow error and my gradients and loss become NaNs.
I tried to run this part of the model with autocast(enabled=False) also trying in combination with things like @custom_fwd(cast_inputs=torch.float32) but haven’t been able to get it to work. If I try these options I get a cudNN error about not finding a kernel to perform convolutions with (in other parts of the network), or I kept getting the under/overflow errors.
torch.Size([32, 6, 75, 20]) torch.Size([32, 1, 1, 20]) are the shapes of scores and mask respectively
Otherwise the mixed precision was working well for me now, as I had to lower the batch_size when turning it off because of memory requirements
Thanks.