(encoder_model, contrast_model, data, optimizer)
| 52 | |
| 53 | |
| 54 | def train(encoder_model, contrast_model, data, optimizer): |
| 55 | encoder_model.train() |
| 56 | optimizer.zero_grad() |
| 57 | z, g, zn = encoder_model(data.x, data.edge_index) |
| 58 | loss = contrast_model(h=z, g=g, hn=zn) |
| 59 | loss.backward() |
| 60 | optimizer.step() |
| 61 | return loss.item() |
| 62 | |
| 63 | |
| 64 | def test(encoder_model, data): |