def __init__(self, output_channel, n_sub_layer,
readout_mode='sum'):
super(WeaveModuleBN, self).__init__()
with self.init_scope():
self.atom_layer = LinearLayer(output_channel, n_sub_layer)
self.pair_layer = LinearLayer(output_channel, n_sub_layer)
self.atom_to_atom = LinearLayer(output_channel, n_sub_layer)
self.pair_to_pair = LinearLayer(output_channel, n_sub_layer)
self.atom_to_pair = AtomToPair(output_channel, n_sub_layer)
self.pair_to_atom = PairToAtom(output_channel, n_sub_layer,
mode=readout_mode)
self.bn = GraphBatchNormalization(output_channel)
self.n_channel = output_channel
self.readout_mode = readout_mode
def __init__(self, *args, **kwargs):
super(SchNetUpdateBN, self).__init__(*args, **kwargs)
with self.init_scope():
self.bn = GraphBatchNormalization(args[0])
