Error in defining the inputs for a class in Tensorflow

I tried giving the input as input channel and output channel as tf.constant tensor. But the code shows error as input for the class must be a vector of {int32,int64}, got shape [2,1]. I defined the GTLayer class as follows :

Importing libraries:

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np

GTLayer class:

class GTLayer(keras.layers.Layer):
    
    def __init__(self, in_channels, out_channels, first=True):
        super(GTLayer, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.first = first
        if self.first == True:
            self.conv1 = GTConv(in_channels, out_channels)
            self.conv2 = GTConv(in_channels, out_channels)
        else:
            self.conv1 = GTConv(in_channels, out_channels)
    
    def forward(self, A, H_=None):
        if self.first == True:
            a = self.conv1(A)
            b = self.conv2(A)
            #H = torch.bmm(a,b)
            H = tf.matmul(a, b)

            #W = [(F.softmax(self.conv1.weight, dim=1)).detach(),(F.softmax(self.conv2.weight, dim=1)).detach()]
            W = [tf.stop_gradient(tf.nn.softmax(self.conv1.weight, axis=1).numpy()),
                 tf.stop_gradient(tf.nn.softmax(self.conv1.weight, axis=1).numpy()) ]

        else:
            a = self.conv1(A)
            #H = torch.bmm(H_,a)
            H = tf.matmul(H_, a)
            #W = [(F.softmax(self.conv1.weight, dim=1)).detach()]
            W = [tf.stop_gradient(tf.nn.softmax(self.conv1.weight, axis=1).numpy())]
        return H,W

GTConv layer

class GTConv(keras.layers.Layer):
    
    def __init__(self, in_channels, out_channels):
        super(GTConv, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        w_init = tf.random_normal_initializer()
        self.weight = tf.Variable(
            initial_value=w_init(shape=(in_channels, out_channels)),
            trainable=True)
        self.bias = None
        self.scale = tf.Variable([0.1] , trainable=False)
        self.reset_parameters()

        #self.weight = nn.Parameter(torch.Tensor(out_channels,in_channels,1,1))
        #self.bias = None
        #self.scale = nn.Parameter(torch.Tensor([0.1]), requires_grad=False)
        
    def reset_parameters(self):
        n = self.in_channels
        tf.fill(self.weight, 9)
        #if self.bias is not None:
        #    fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
        #   bound = 1 / math.sqrt(fan_in)
        #    nn.init.uniform_(self.bias, -bound, bound)
    
    def forward(self, A):
        A = tf.add_n(tf.nn.softmax(self.weight))
        return A 

Error:

input:

inp = tf.constant([4])
out = tf.constant([2])
d = GTLayer(inp, out)

--------------------------------------------------------------------------- InvalidArgumentError Traceback (most recent call last) in () ----> 1 d = GTLayer(inp, out)

5 frames /usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/ops.py in raise_from_not_ok_status(e, name) 7184 def raise_from_not_ok_status(e, name): 7185 e.message += (" name: " + name if name is not None else “”) → 7186 raise core._status_to_exception(e) from None # pylint: disable=protected-access 7187 7188

InvalidArgumentError: shape must be a vector of {int32,int64}, got shape [2,1] [Op:RandomStandardNormal]