()
| 53 | |
| 54 | @partial(mx.compile, inputs=state, outputs=state) |
| 55 | def step(): |
| 56 | loss_and_grad_fn = nn.value_and_grad(gcn, forward_fn) |
| 57 | (loss, y_hat), grads = loss_and_grad_fn( |
| 58 | gcn, x, adj, y, train_mask, args.weight_decay |
| 59 | ) |
| 60 | optimizer.update(gcn, grads) |
| 61 | return loss, y_hat |
| 62 | |
| 63 | best_val_loss = float("inf") |
| 64 | cnt = 0 |