Keras Custom Layer: __init__takes 1 positional argument but 2 were given

Question

I'm trying to create a layer that will split the columns of a [3,5] tensor into a [3,2] and [3,3] tensor respectively. For example,

[[0., 1., 0.,  0., 0.],
 [1., 0., -1., 0., 0.],
 [1., 0., 1.,  1., 0.]]

into,

[[0., 1.],
 [1., 0.],
 [1., 0.]]

[[0.,  0., 0.],
 [-1., 0., 0.],
 [1.,  1., 0.]]

This is the custom layer I've tried to build,

import tensorflow as tf
from tensorflow.keras.layers import Layer

class NodeFeatureSplitter(Layer):
    def __init__(self):
        super(NodeFeatureSplitter, self).__init__()
        
    def call(self, x):
        h_feat = x[...,:2]
        x_feat = x[...,-3:]
        
        return h_feat, x_feat

However when I call this layer on the following example I get the aforementioned error,

x = tf.constant([[0., 1., 0., 0., 0.],[1., 0., -1., 0., 0.],[1., 0., 1., 1., 0.]])

h_feat, x_feat = NodeFeatureSplitter(x)
print(h_feat)
print(x_feat)

TypeError: __ init __() takes 1 positional argument but 2 were given

Can someone highlight what I'm doing wrong?

Thanks <3