model.predict does not work with stochastic output layers

`model.predict` doesn't work for non-Tensor outputs, including Tensor-convertible objects like `ed.RandomVariable`.

For now, the workaround is to replace `model.predict` as below with an explicit for loop over the data.

```python
dataset_test = dataset_test.repeat().batch(batch_size)
test_steps = ds_info.splits['test'].num_examples // batch_size

predictions = model.predict(dataset_test, verbose=1, steps=test_steps) # raises error
logits = predictions.distribution.logits # predicted logits of full dataset
```

```python
dataset_test = dataset_test.batch(batch_size)

logits = []
for features, _ in dataset_test:
predictions = model(features)
logits.append(predictions.distribution.logits)

logits = tf.concat(logits, axis=0) # predicted logits of full dataset
```

Note to loop over tf data, you need to use TF 2.0 behavior; otherwise you need to use a tf.Session with the deprecated iterator design.