(self, s_i, v_i, s_j, v_j, edge_attr)
| 283 | return _split(message, self.vo) |
| 284 | |
| 285 | def message(self, s_i, v_i, s_j, v_j, edge_attr): |
| 286 | v_j = v_j.view(v_j.shape[0], v_j.shape[1]//3, 3) |
| 287 | v_i = v_i.view(v_i.shape[0], v_i.shape[1]//3, 3) |
| 288 | message = tuple_cat((s_j, v_j), edge_attr, (s_i, v_i)) |
| 289 | message = self.message_func(message) |
| 290 | return _merge(*message) |
| 291 | |
| 292 | |
| 293 | class GVPConvLayer(nn.Module): |