Unverified Commit 305aea94 authored by Karl Leswing's avatar Karl Leswing Committed by GitHub
Browse files

Merge pull request #1178 from souravsingh/fix-cpu 

Remove CPU pin from GraphGather
parents 6795a943 0ce6d470

deepchem/models/tensorgraph/layers.py

+30 −31
Original line number Diff line number Diff line
@@ -2489,7 +2489,6 @@ class GraphGather(Layer): 
    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...]

CRA Git | Maintained and supported by SUSTech CRA and CCSE