(encoder_model, contrast_model, data, optimizer)
| 60 | |
| 61 | |
| 62 | def train(encoder_model, contrast_model, data, optimizer): |
| 63 | encoder_model.train() |
| 64 | optimizer.zero_grad() |
| 65 | z1, z2, g1, g2, z1n, z2n = encoder_model(data.x, data.edge_index) |
| 66 | loss = contrast_model(h1=z1, h2=z2, g1=g1, g2=g2, h1n=z1n, h2n=z2n) |
| 67 | loss.backward() |
| 68 | optimizer.step() |
| 69 | return loss.item() |
| 70 | |
| 71 | |
| 72 | def test(encoder_model, data): |