Implementing a LoRA wrapper for Conv2D in Tensorflow

def call(self, inputs, training=None):
    # Compute LoRA:
    deltaW = tf.matmul(self.loraB, self.loraA)
    deltaW = tf.reshape(deltaW, self.kernelshape)
    # Apply scaling based on training:
    if training and self.loraenabled and not self.merged:
        effectivedelta = self.scaling * deltaW
    else:
        effectivedelta = tf.zeroslike(deltaW)
    outputs = tf.nn.conv2d(
        inputs,
        self.originallayer.kernel + effectivedelta,
        strides=self.originallayer.strides,
        padding=self.originallayer.padding.upper()
    )
    if self.originallayer.usebias:
        outputs = tf.nn.biasadd(outputs, self.originallayer.bias)
    if self.originallayer.activation is not None:
        outputs = self.originallayer.activation(outputs)
    return outputs

def build(self, inputshape):
    # ... existing code ...
    # Always trainable, control via scaling
    self.loraA = self.addweight(
        name="LoRAmatA",
        shape=loraAshape,
        initializer=keras.initializers.HeUniform(),
        trainable=True,  # Always True
    )
    self.loraB = self.addweight(
        name="LoRAmatB",
        shape=loraBshape,
        initializer=keras.initializers.Zeros(),
        trainable=True,  # Always True
    )
    # Control via this flag
    self.loraactive = tf.Variable(
        initialvalue=self.loraenabled,
        trainable=False,
        dtype=tf.bool
    )
def call(self, inputs):
    deltaW = tf.matmul(self.loraB, self.loraA)
    deltaW = tf.reshape(deltaW, self.kernelshape)
    # Use tf.cond for runtime switching
    effectivescaling = tf.cond(
        self.loraactive,
        lambda: self.scaling,
        lambda: 0.0
    )
    #...