Unable to use heteroscedastic layers in tf.keras 

`TypeError: Dimension value must be integer or None or have an __index__ method, got value '<attribute 'shape' of 'numpy.generic' objects>' with type '<class 'getset_descriptor'>'`

```
hidden_layer = tf.keras.layers.Dense(16)
inputs = tf.keras.Input((224,224,3))
out = hidden_layer(inputs)
output_layer = ed.layers.MCSoftmaxDense(num_classes=2)
output_layer(out)
```