[BUG] predict_on_batch fails after calling model_wo_softmax

[lopf33]

After calling innvestigate.model_wo_softmax using predict_on_batch on the model fails with: AttributeError: 'Functional' object has no attribute '_training_endpoints'

Steps to reproduce the bug

import tensorflow as tf
from tensorflow import keras
import innvestigate
tf.compat.v1.disable_eager_execution()

model = tf.keras.models.load_model('your_model_path')
model_wo_softmax = innvestigate.model_wo_softmax(model)
model_wo_softmax.predict_on_batch(dataset.batch(32))

When using model.layers[-1].activation = None to remove the softmax predict_on_batch does not fail when called directly but still fails in the analyzer.analyze call.

[lopf33]

The issue was caused by me trying to use tf.data.Datasets with eager execution disabled. Managing the data manually in a list fixed the issue.