def forward(self, x, identity=None):
        """Forward function for `FFN`.
        The function would add x to the output tensor if residue is None.
        """
        
        if self.with_cp and x.requires_grad:
            out = cp.checkpoint(self.layers, x)
        else:
            out = self.layers(x)
            
        if not self.add_identity:
            return self.dropout_layer(out)
        if identity is None:
            identity = x
        return identity + self.dropout_layer(out)
Forward function for `FFN`.
        The function would add x to the output tensor if residue is None.
        
EVA-01/seg/mmseg_custom/models/utils/transformer.py