(encoder_model, contrast_model, data, optimizer)
| 101 | |
| 102 | |
| 103 | def train(encoder_model, contrast_model, data, optimizer): |
| 104 | encoder_model.train() |
| 105 | optimizer.zero_grad() |
| 106 | _, _, h1_pred, h2_pred, h1_target, h2_target = encoder_model(data.x, data.edge_index, data.edge_attr) |
| 107 | loss = contrast_model(h1_pred=h1_pred, h2_pred=h2_pred, h1_target=h1_target.detach(), h2_target=h2_target.detach()) |
| 108 | loss.backward() |
| 109 | optimizer.step() |
| 110 | encoder_model.update_target_encoder(0.99) |
| 111 | return loss.item() |
| 112 | |
| 113 | |
| 114 | def test(encoder_model, data): |
no test coverage detected