tensor.bmm() with unexpected autocast to float16

class StyleLoss(nn.Module):
    r"""
    Perceptual loss, VGG-based
    https://arxiv.org/abs/1603.08155
    https://github.com/dxyang/StyleTransfer/blob/master/utils.py
    """

    def __init__(self):
        super(StyleLoss, self).__init__()
        self.add_module('vgg', VGG19())
        self.criterion = torch.nn.L1Loss()

    @amp.float_function
    def compute_gram(self, x):
        b, ch, h, w = x.size()
        f = x.view(b, ch, w * h) # fp32
        f_T = f.transpose(1, 2) # fp32
        G = f.bmm(f_T) / (h * w * ch) # fp16

        return G

    @amp.float_function
    def __call__(self, x, y):
        # Compute features
        x_vgg, y_vgg = self.vgg(x), self.vgg(y)

        # Compute loss
        style_loss = 0.0
        style_loss += self.criterion(self.compute_gram(x_vgg['relu2_2']), self.compute_gram(y_vgg['relu2_2']))
        style_loss += self.criterion(self.compute_gram(x_vgg['relu3_4']), self.compute_gram(y_vgg['relu3_4']))
        style_loss += self.criterion(self.compute_gram(x_vgg['relu4_4']), self.compute_gram(y_vgg['relu4_4']))
        style_loss += self.criterion(self.compute_gram(x_vgg['relu5_2']), self.compute_gram(y_vgg['relu5_2']))

        return style_loss

G = f.bmm(f_T) / (h * w * ch) this operation will cause nan with float16. Although f and f_T are both float32, f.bmm(f_T) will also be autocasted to float16 and cause nan. How to deal with it?

Start with the shown StyleLoss.compute_gram example, especially the f.bmm(f_T) operation and its @amp.float_function context. Reproduce the reported NaN behavior and trace the autocast handling; done means the operation no longer unexpectedly produces NaN when both inputs are float32.