Get shapes right

from deepchem.models.tensorgraph.layers import Flatten, Dense, SoftMax, \

  Variable, \

  Feature, Layer, Add, BatchNorm, Conv2D

  Feature, Layer, Add, BatchNorm, Conv2D, Squeeze

from deepchem.rl.a3c import _Worker

    d4 = BatchNorm(in_layers=[d4])

    d5 = Dense(in_layers=[d4], activation_fn=None, out_channels=9)

    value = Dense(in_layers=[d4], activation_fn=None, out_channels=1)

    value = Squeeze(squeeze_dims=1, in_layers=[value])

    probs = SoftMax(in_layers=[d5])

    return {'action_prob': probs, 'value': value}

    return out_tensor

class Squeeze(Layer):

  def __init__(self, squeeze_dims, **kwargs):

    self.squeeze_dims = squeeze_dims

    super(Squeeze, self).__init__(**kwargs)

  def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):

    inputs = self._get_input_tensors(in_layers)

    parent_tensor = inputs[0]

    out_tensor = tf.squeeze(parent_tensor, squeeze_dims=self.squeeze_dims)

    if set_tensors:

      self.out_tensor = out_tensor

    return out_tensor

class Transpose(Layer):

  def __init__(self, perm, **kwargs):

      action = np.random.choice(np.arange(n_actions), p=probabilities[0])

      actions.append(np.zeros(n_actions))

      actions[i][action] = 1.0

      values.append(value[0])

      values.append(value)

      rewards.append(self.env.step(action))

    if not self.env.terminated:

      # Add an estimate of the reward for the rest of the episode.