Stateful LSTM in custom training loop

I am writing a simple custom model + training in tensorflow. My goal is to build a stateful LSTM based model and being able to reset the states when I want to.

So far this is my custom model:

class ResNetModel(Model):
    def __init__(self, num_inputs, **kwargs):
        """
        The class initialiser should call the base class initialiser, passing any keyword
        arguments along. It should also create the layers of the network according to the
        above specification.
        """
        super(ResNetModel, self).__init__(**kwargs)
        self.lstm_1 = tf.keras.layers.LSTM(units=32, input_shape=(None, num_inputs), return_sequences=True)
        self.dense = tf.keras.layers.Dense(units=1, activation=None)

    def call(self, inputs, training=False):
        """
        This method should contain the code for calling the layer according to the above
        specification, using the layer objects set up in the initialiser.
        """
        x = self.lstm_1(inputs)
        y = self.dense(x)
        return y + inputs

And this is my custom training loop (I am omitting the whole code because it is quite big, but the function is self contained for the purpose of my question):

def run_training(self, in_train, out_train, epoch_loss, epoch_error, n_skip, n_block):
    n_samples = in_train.shape[1]

    self.model.reset_states()  # clear existing state
    self.model(in_train[:, :n_skip, :])  # process some samples to build up state

    for n in range(n_skip, n_samples - n_block, n_block):

        # compute loss
        with tf.GradientTape() as tape:
            y_pred = self.model(in_train[:, n:n + n_block, :])
            loss = self.loss_func(out_train[:, n:n + n_block, :], y_pred)

        grads = tape.gradient(loss, self.model.trainable_variables)


        self.opt.apply_gradients(zip(grads, self.model.trainable_variables))

        epoch_loss.update_state(loss)

        epoch_error.update_state(out_train[:, n:n + n_block, :], y_pred)

And it trains fine, the whole code works as expected.

Then I make predictions like this:

for i in range(0, math.floor(24000/4096)):
    predictions[i*4096: (i+1)*4096] = np.array(residual_net.model(X_test[idx][i*4096: (i+1)*4096].reshape(1, 4096, 1))).ravel()

So basically I am passing my input test to my model in residual_net.model(my_test_data) (the numpy slicing ecc... is just to make my input data coherent with the network, it works fine).

However, when I make predictions with my trained network (to give some context, it is working with audio data) I have the output audio that is as expected (an input song processed by the network that adds some distortion), but there are clicks in the output audio that are directly related to the input buffer size.

To make this point clearer: if I predict on chuncks of 512 samples, I have clicks every 512 samples, if I predict every 4096 samples, I have clicks every 4096.

This behaviour is pretty similar to the one you have with IIR filters that do not carry its filter state across audio buffers, so that got me thinking that my LSTM network is not working stateful as I expected.

So my question is:

does Tensorflow automatically reset the state of the network after each processed buffer (even in the case of custom training-prediction loops) if the parameter stateful = True is not specified in the LSTM layer? I found no information about this, but I expected that behaviour for "standard" training (.fit/.predict functions) and not for custom training loops. Does this hold also for the training step? (so basically I am messing up also the training)