Error when using WeaveModel with gaussian_expand option

Steps to reproduce:

import numpy as np
import tensorflow as tf
featurizer = dc.feat.WeaveFeaturizer()
X = featurizer(["C", "CC"])
y = np.array([[1], [0]])
dataset = dc.data.NumpyDataset(X, y)

model = dc.models.WeaveModel(n_tasks=1, n_weave=2, fully_connected_layer_sizes=[100], mode="regression", nb_epoch = 1,
                             activation=tf.nn.sigmoid, final_conv_activation_fn=tf.nn.sigmoid, compress_post_gaussian_expansion=True, 
                             gaussian_expand=False)
loss = model.fit(dataset)

model.predict(dataset)

The above code raises the following error:

ValueError: Exception encountered when calling layer "weave_gather_23" (type WeaveGather).

in user code:

    File "/usr/local/lib/python3.7/dist-packages/deepchem/models/layers.py", line 3086, in call  *
        output_molecules = tf.matmul(output_molecules, self.W) + self.b

    ValueError: Dimensions must be equal, but are 128 and 1408 for '{{node weave_gather_23/MatMul}} = MatMul[T=DT_FLOAT, transpose_a=false, transpose_b=false](weave_gather_23/SegmentSum, weave_gather_23/MatMul/ReadVariableOp)' with input shapes: [?,128], [1408,128].

This is probably due to the use of gaussian_expand and compress_post_gaussian_expansion options. It occurs only when they are (True, False).