(model, x, y)
| 58 | |
| 59 | def train_epoch(model, train_iter, optimizer, epoch): |
| 60 | def train_step(model, x, y): |
| 61 | output = model(x) |
| 62 | loss = mx.mean(nn.losses.cross_entropy(output, y)) |
| 63 | acc = mx.mean(mx.argmax(output, axis=1) == y) |
| 64 | return loss, acc |
| 65 | |
| 66 | state = [model.state, optimizer.state] |
| 67 |
nothing calls this directly
no outgoing calls
no test coverage detected