Seems like the function predict_batch, formerly predict_macroscopic is taking a Tensorflow dataset, from which only the first element ({'inputs': self.input_data, 'params': param_data}) is used. The "inputs" entry of that dictionary contains num_samples copies of the same input sequence, which could be optimized.