1 files changed, 16 insertions, 0 deletions
diff --git a/include/train.h b/include/train.h
new file mode 100644
index 0000000..9cf613e
--- /dev/null
+++ b/include/train.h
@@ -0,0 +1,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);