Gradient overflows when self-attention module added

When I add self-attention module to the network, gradient overflows and loss becomes Nan. (opt_level=O2)
Following is code for the self-attention module (got from link) I used,

class CGDBlock3D(nn.Module):
    def __init__(self, in_channels):
        super(CGDBlock3D, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool3d(1)
        self.max_pool = nn.AdaptiveMaxPool3d(1)
        self.softmax = nn.Softmax(dim=1)

        self.w0 = nn.Parameter(torch.ones(in_channels, 1))
        self.w1 = nn.Parameter(torch.ones(in_channels, 1))
        self.w2 = nn.Parameter(torch.ones(in_channels, 1))

        self.bias0 = nn.Parameter(torch.zeros(1, in_channels, 1, 1, 1))
        self.bias1 = nn.Parameter(torch.zeros(1, in_channels, 1, 1, 1))
        self.bias2 = nn.Parameter(torch.zeros(1, in_channels, 1, 1, 1))

        nn.init.xavier_uniform_(self.w0)
        nn.init.xavier_uniform_(self.w1)
        nn.init.xavier_uniform_(self.w2)

    def cgd(self, x, N, C):
        g = self.avg_pool(x).view(N, C, 1, 1, 1)
        f = self.max_pool(x).view(N, C, 1, 1, 1)

        g_s = self.softmax(g)  # b ,c ,1 ,1, 1

        psi = torch.matmul(g.view(N, C), self.w0).view(N, 1, 1, 1, 1)
        phi = torch.matmul(f.view(N, C), self.w1).view(N, 1, 1, 1, 1)

        psi_s = torch.tanh(psi * g_s + self.bias0)  # b ,c ,1 ,1, 1
        phi_s = torch.tanh(phi * g_s + self.bias1)  # b ,c ,1 ,1, 1

        gf = torch.matmul(phi_s.view(N, C), self.w2).view(N, 1, 1, 1, 1)
        gf_tanh = torch.tanh(gf * psi_s + self.bias2).view(N, C, 1, 1, 1)

        z = x * (1 + gf_tanh)
        return z

    def forward(self, x):
        N, C, T, _, _ = x.size()
        out = self.cgd(x, N, C)
        return out

I'm not sure which part of the module makes the problem.
I assumed softmax but it happened even I remove the softmax from the code.

Start by reproducing the reported self-attention module with Apex mixed precision at opt_level=O2, using the provided CGDBlock3D code. Compare runs with and without the softmax and isolate the operation that first produces overflow or NaN. Done means identifying the failing operation and documenting a minimal reproduction or actionable diagnosis.