    def test_ffn(self):
        """
        Test FFN block.
        """
        x = torch.randn(1, 3, 1, 3)
        layer1 = FFNBlock(in_channels=3)
        layer2 = FFNBlock(in_channels=3, hidden_channels=2, out_channels=2)
        y1 = layer1(x)
        y2 = layer2(x)
        # print(y1.shape, y2.shape)
        self.assertTrue(y1.shape == torch.Size([1, 3, 1, 3]))
        self.assertTrue(y2.shape == torch.Size([1, 2, 1, 3]))
Test FFN block.
tests/unittests/models/layers/test_ffn.py