Pooling architecture using the TW layer.
(self, n_features, n_templates, n_template_nodes, pooling_layer)
| 100 | """ |
| 101 | |
| 102 | def __init__(self, n_features, n_templates, n_template_nodes, pooling_layer): |
| 103 | """ |
| 104 | Pooling architecture using the TW layer. |
| 105 | """ |
| 106 | super().__init__() |
| 107 | |
| 108 | self.n_features = n_features |
| 109 | self.n_templates = n_templates |
| 110 | self.n_template_nodes = n_template_nodes |
| 111 | |
| 112 | self.TFGW = pooling_layer |
| 113 | |
| 114 | self.linear = Linear(self.n_templates, 1) |
| 115 | |
| 116 | def forward(self, x, edge_index, batch=None): |
| 117 | x = self.TFGW(x, edge_index, batch=batch) |