include/train.h


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

#pragma once

#include "dynmlp.h"
#include "adam.h"

typedef struct {
    double loss;
    int nfe_fwd;
    int nfe_bwd;
} TrainStepResult;

TrainStepResult train_step(const DynMLP *net, double *theta,
                           const double **z0s, const double **targets,
                           double t0, double t1, int batch_size,
                           Adam *adam, double atol, double rtol,
                           int num_checkpoints);