row = torch.arange(batch.size(0), device=batch.device)
edge_index = torch.stack([row, batch], dim=0)
out = global_add_pool(xs[-1], batch)
for conv, gru, x in zip(self.mol_convs, self.mol_grus, xs):
h = F.elu_(conv((x, out), edge_index))
out = gru(h, out)
After Change
x = F.leaky_relu_(self.lin1(x))
h = F.elu_(self.atom_convs[0](x, edge_index, edge_attr))
h = F.dropout(h, p=self.dropout, training=self.training)
x = self.atom_grus[0](h, x).relu_()
for conv, gru in zip(self.atom_convs[1:], self.atom_grus[1:]):
h = F.elu_(conv(x, edge_index))
h = F.dropout(h, p=self.dropout, training=self.training)
x = gru(h, x).relu_()
// Molecule Embedding:
row = torch.arange(batch.size(0), device=batch.device)
edge_index = torch.stack([row, batch], dim=0)
out = global_add_pool(x, batch).relu_()
for t in range(self.num_timesteps):
h = F.elu_(self.mol_conv((x, out), edge_index))
h = F.dropout(h, p=self.dropout, training=self.training)
out = self.mol_gru(h, out).relu_()
// Predictor:
out = F.dropout(out, p=self.dropout, training=self.training)
return self.lin2(out)
