Merge pull request #1081 from lilleswing/pin-graphgather

    if len(inputs) != 2:

      raise ValueError()

    labels, logits = inputs[0], inputs[1]

    out_tensor = tf.nn.softmax_cross_entropy_with_logits(

    out_tensor = tf.nn.softmax_cross_entropy_with_logits_v2(

        logits=logits, labels=labels)

    if set_tensors:

      self.out_tensor = out_tensor

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

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

    with tf.device('/cpu'):

      inputs = self._get_input_tensors(in_layers)

      # x = [atom_features, deg_slice, membership, deg_adj_list placeholders...]

      # Sum over atoms for each molecule

      sparse_reps = [

        tf.reduce_mean(activated, 0, keep_dims=True)

          tf.reduce_mean(activated, 0, keepdims=True)

          for activated in activated_par

      ]

      max_reps = [

        tf.reduce_max(activated, 0, keep_dims=True)

          tf.reduce_max(activated, 0, keepdims=True)

          for activated in activated_par

      ]