All values input SLA and the resulting tensor become NaN

[ZouYa99]

Hello, I want to incorporate SLA into the video VAE network architecture. I reshape the 5D video tensor and input it into the SLA layer, but all values in the resulting tensor become NaN. The input video tensor has values ranging from -1 to 1, and I'm using bf16 precision. The SLA wrapper class I'm using is as follows. I've confirmed that removing the SLA layer results in normal values. Do you have any suggestions?

class SLABlock(nn.Module):
    def __init__(
        self, 
        channels, 
        head_dim=128,
        topk=0.2,  
        feature_map="relu", 
        block_size=64
    ):
        super().__init__()
        self.channels = channels
        self.head_dim = head_dim
        self.num_heads = channels // head_dim
        
        self.attn = SparseLinearAttention(
            head_dim=self.head_dim,
            topk=topk,
            feature_map=feature_map,
            BLKQ=block_size,
            BLKK=block_size,
        )

    def forward(self, x):
        """
        Input:  [Batch, Channels, Time, Height, Width]
        Output: [Batch, Channels, Time, Height, Width]
        """
        b, c, t, h, w = x.shape

        x = x.permute(0, 2, 3, 4, 1).flatten(1, 3) 
        
        x = x.view(b, -1, self.num_heads, self.head_dim)
        x = x.transpose(1, 2).contiguous()
        
        out = self.attn(x, x, x) #out become NaN

        out = out.transpose(1, 2).reshape(b, t * h * w, c)
        out = out.view(b, t, h, w, c).permute(0, 4, 1, 2, 3)
        
        return out