Is there any possibility to use Custom layer trainable weight as a parameter for tensorflow addons function rotation?

What I want to do is learn my custom layer what is the best rotation for images of one domain. Im using tensorflow addons functions and want to pass weight value as parameter. There is my custom layer

class RotationLayer(Layer):
def __init__(self, num_outputs, name = 'RotationLayer', **kwargs):
    super(RotationLayer, self).__init__(**kwargs)
    self.num_outputs = num_outputs

def build(self, input_shape):
    super(RotationLayer, self).build(input_shape)

    self.w = self.add_weight(
        shape=[1], initializer="uniform", trainable=True
    )

def call(self, inputs, **kwargs):
    rotated = tfa.image.rotate(inputs, (self.w*100)*math.pi/180)
    rotated = tf.reshape(rotated, shape=[-1,64,64,1])
    return rotated

def compute_output_shape(self, input_shape):
    return (input_shape[0], self.output_dim)

def get_config(self):
    base = super(RotationLayer, self).get_config()
    base['num_outputs'] = self.num_outputs
    return dict(list(base.items()))

This layer works but totaly ruin my CNN so I want to ask if is there any other possibility. My goal is to measure impact of learned data augmentation layer on CNN architecture a metrics.

Neural network architecture

model = Sequential()
#model.add(flipLayer)
#model.add(brightnessLayer)
#model.add(rotatioLayer)

model.add(Conv2D(48, 5, padding = 'same', activation = 'relu', input_shape=(64,64,1)))
model.add(rotatioLayer)
model.add(MaxPool2D(pool_size = 3, strides = 2))
model.add(Conv2D(64, 5, padding = 'same', activation = 'relu'))
model.add(MaxPool2D(pool_size = 3, strides = (2,2)))
model.add(Conv2D(32, 3, padding = 'same', activation = 'relu'))
model.add(Dropout(0.1))
model.add(Dense(1024, activation = 'relu'))
model.add(Flatten())
model.add(Dense(1, activation = 'sigmoid'))
model.compile(loss = 'binary_crossentropy', optimizer = 'adam', metrics = ['acc'], run_eagerly=True)
model.summary()