(encoder_model, contrast_model, data, optimizer)
| 54 | |
| 55 | |
| 56 | def train(encoder_model, contrast_model, data, optimizer): |
| 57 | encoder_model.train() |
| 58 | optimizer.zero_grad() |
| 59 | z, z1, z2 = encoder_model(data.x, data.edge_index, data.edge_attr) |
| 60 | h1, h2 = [encoder_model.project(x) for x in [z1, z2]] |
| 61 | loss = contrast_model(h1, h2) |
| 62 | loss.backward() |
| 63 | optimizer.step() |
| 64 | return loss.item() |
| 65 | |
| 66 | |
| 67 | def test(encoder_model, data): |