diff --git a/examples/template_training.ipynb b/examples/template_training.ipynb
new file mode 100644
index 0000000..d11b53a
--- /dev/null
+++ b/examples/template_training.ipynb
@@ -0,0 +1,510 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "d31d3171",
+   "metadata": {},
+   "source": [
+    "# Template for autoencoder prototyping"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d67dad03-9d76-4891-82ff-7e19d1369a24",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import numpy as np\n",
+    "from matplotlib.colors import LinearSegmentedColormap\n",
+    "import matplotlib.pyplot as plt\n",
+    "import matplotlib.animation as animation\n",
+    "%matplotlib inline\n",
+    "\n",
+    "# Name for the case\n",
+    "name = 'template'"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fceceb92",
+   "metadata": {},
+   "source": [
+    "## Hyper-parameters"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6b05f7c3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dt = 1.0e-5\n",
+    "ae_weight = 1e0\n",
+    "sindy_weight = 1.e-6\n",
+    "coef_weight= 1.e-9\n",
+    "lr = 1e-5"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "55bc1410",
+   "metadata": {},
+   "source": [
+    "## Training data\n",
+    "\n",
+    "- Consider we only have one time-series simulation.\n",
+    "- suppose we have grid of size (20x30) points.\n",
+    "- solution is a 2-dimensional vector (on each grid point)\n",
+    "- A simulation of 50 time steps.\n",
+    "- the first index of snapshot array should be timestep.\n",
+    "- we consider dof as solution dimension, and number of particles as grid size.\n",
+    "\n",
+    "Can load other data files for your custom autoencoder."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ca3062c5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ntrain = 1\n",
+    "sol_dim = 2\n",
+    "grid_dim = [20,30]\n",
+    "nt = 50\n",
+    "\n",
+    "# number of cases x number of time step x solution dimension x grid size\n",
+    "snapshots = np.random.rand(ntrain, nt, sol_dim, grid_dim[0], grid_dim[1])\n",
+    "X_train = torch.Tensor(snapshots)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1b6241c5",
+   "metadata": {},
+   "source": [
+    "## Setting up a physics wrapper.\n",
+    "\n",
+    "You can choose from `lasdi.physics` module, or create a custom physics model. Following snippet is an example with some necessary specifications."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3a0349ce",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from lasdi.physics import Physics\n",
+    "\n",
+    "class CustomPhysicsModel(Physics):\n",
+    "    def __init__(self):\n",
+    "        self.dim = 2\n",
+    "        self.nt = nt\n",
+    "        self.dt = dt\n",
+    "        self.qdim = sol_dim\n",
+    "        self.grid_size = grid_dim\n",
+    "        self.qgrid_size = [sol_dim] + grid_dim\n",
+    "\n",
+    "        ''' Set up a grid as you see fit. '''\n",
+    "        self.x_grid = np.zeros([self.dim] + self.grid_size)\n",
+    "        xg = np.linspace(0., 1., grid_dim[1])\n",
+    "        yg = np.linspace(0., 1., grid_dim[0])\n",
+    "        self.x_grid[0], self.x_grid[1] = np.meshgrid(xg, yg)\n",
+    "        return\n",
+    "    \n",
+    "    ''' See lasdi.physics.Physics class for necessary subroutines '''\n",
+    "\n",
+    "physics = CustomPhysicsModel()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9996b792",
+   "metadata": {},
+   "source": [
+    "## Setting up autoencoder\n",
+    "\n",
+    "You can choose from `lasdi.latent_space` module, or create a custom autoencoder. The snippet below uses a built-in `Autoencoder` class."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c53127b3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from lasdi.latent_space import Autoencoder\n",
+    "\n",
+    "n_train = 1\n",
+    "hidden_units = [3000, 300, 300, 100, 100, 30]\n",
+    "# latent space dimension\n",
+    "n_z = 5\n",
+    "print(physics.grid_size, hidden_units, n_z)\n",
+    "\n",
+    "# I think these options are straightforward.\n",
+    "ae_cfg = {'hidden_units': hidden_units, 'latent_dimension': n_z, 'activation': 'softplus'}\n",
+    "\n",
+    "ae = Autoencoder(physics, ae_cfg)\n",
+    "best_loss = np.Inf"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a0ea7a53",
+   "metadata": {},
+   "source": [
+    "## Setting up latent dynamics\n",
+    "\n",
+    "Similarly, we either choose a latent dynamics model from `lasdi.latent_dynamics` module or create a custom model. Here we use `SINDy` class."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ad4ae3ce",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from lasdi.latent_dynamics.sindy import SINDy\n",
+    "\n",
+    "sindy_options = {'sindy': {'fd_type': 'sbp12'}} # finite-difference operator for computing time derivative of latent trajectory.\n",
+    "ld = SINDy(ae.n_z, physics.nt, sindy_options)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b5d98899",
+   "metadata": {},
+   "source": [
+    "## Training: (1) Running a custom optimization process"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "36bacfa9",
+   "metadata": {},
+   "source": [
+    "Set up optimizer and loss term"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "31de72b7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "optimizer = torch.optim.Adam(ae.parameters(), lr = lr)\n",
+    "MSE = torch.nn.MSELoss()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "55cfb184",
+   "metadata": {},
+   "source": [
+    "Set up device: `'cuda'`, `'mps'` (apple) or `'cpu'`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "35cd5db2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "device = 'cpu'\n",
+    "d_ae = ae.to(device)\n",
+    "d_Xtrain = X_train.to(device)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "65bc550a",
+   "metadata": {},
+   "source": [
+    "This is a part of GPLaSDI training procedure, where the training is executed without on-the-fly sampling."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ff367763",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "n_iter = 10\n",
+    "save_interval = 1\n",
+    "hist_file = '%s.loss_history.txt' % name\n",
+    "\n",
+    "loss_hist = np.zeros([n_iter, 4])\n",
+    "grad_hist = np.zeros([n_iter, 4])\n",
+    "for iter in range(n_iter):\n",
+    "    optimizer.zero_grad()\n",
+    "    d_ae = ae.to(device)\n",
+    "    d_Z = d_ae.encoder(d_Xtrain)\n",
+    "    d_Xpred = d_ae.decoder(d_Z)\n",
+    "    Z = d_Z.cpu()\n",
+    "\n",
+    "    loss_ae = MSE(d_Xtrain, d_Xpred)\n",
+    "    coefs, loss_sindy, loss_coef = ld.calibrate(Z, physics.dt, compute_loss=True, numpy=True)\n",
+    "\n",
+    "    max_coef = np.abs(np.array(coefs)).max()\n",
+    "    loss = ae_weight * loss_ae + sindy_weight * loss_sindy / n_train + coef_weight * loss_coef / n_train\n",
+    "\n",
+    "    loss_hist[iter] = [loss.item(), loss_ae.item(), loss_sindy.item(), loss_coef.item()]\n",
+    "\n",
+    "    loss.backward()\n",
+    "\n",
+    "    optimizer.step()\n",
+    "\n",
+    "    if ((loss.item() < best_loss) and (iter % save_interval == 0)):\n",
+    "        torch.save(d_ae.cpu().state_dict(), './%s_checkpoint.pt' % name)\n",
+    "        best_loss = loss.item()\n",
+    "\n",
+    "    # print(\"Iter: %05d/%d, Loss: %.5e\" % (iter + 1, n_iter, loss.item()))\n",
+    "    print(\"Iter: %05d/%d, Loss: %.5e, Loss AE: %.5e, Loss SI: %.5e, Loss COEF: %.5e, max|c|: %.5e\"\n",
+    "            % (iter + 1, n_iter, loss.item(), loss_ae.item(), loss_sindy.item(), loss_coef.item(), max_coef))\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "015d02f5",
+   "metadata": {},
+   "source": [
+    "Save the history and the trained parameters."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fb293bb3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "np.savetxt('%s.loss_history.txt' % name, loss_hist)\n",
+    "\n",
+    "if (loss.item() < best_loss):\n",
+    "    torch.save(d_ae.cpu().state_dict(), './%s_checkpoint.pt' % name)\n",
+    "    best_loss = loss.item()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "97f719d3",
+   "metadata": {},
+   "source": [
+    "Load a checkpoint and re-evaluate the losses."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "00f7da25",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "state_dict = torch.load('%s_checkpoint.pt' % name)\n",
+    "ae.load_state_dict(state_dict)\n",
+    "ae.eval()\n",
+    "d_ae = ae.to(device)\n",
+    "d_Z = d_ae.encoder(d_Xtrain)\n",
+    "d_Xpred = d_ae.decoder(d_Z)\n",
+    "Z = d_Z.cpu()\n",
+    "print(Z.shape)\n",
+    "\n",
+    "loss_ae = MSE(d_Xtrain, d_Xpred)\n",
+    "coefs, loss_sindy, loss_coef = ld.calibrate(Z, physics.dt, compute_loss=True, numpy=True)\n",
+    "\n",
+    "max_coef = np.abs(np.array(coefs)).max()\n",
+    "loss = loss_ae + sindy_weight * loss_sindy / n_train + coef_weight * loss_coef / n_train\n",
+    "print(loss.item())\n",
+    "print(loss_ae.item())\n",
+    "print((sindy_weight * loss_sindy / n_train).item())\n",
+    "print((coef_weight * loss_coef / n_train).item())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6516cfa7",
+   "metadata": {},
+   "source": [
+    "Visualize the loss history"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d62570fc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "loss_hist = np.loadtxt('%s.loss_history.txt' % name)\n",
+    "\n",
+    "plt.figure(1)\n",
+    "plt.loglog(loss_hist[:, 0])\n",
+    "plt.xlabel('iteration')\n",
+    "plt.ylabel('loss')\n",
+    "\n",
+    "plt.rcParams.update({'font.size': 15})\n",
+    "plt.figure(2, figsize=(8, 6))\n",
+    "plt.loglog(loss_hist[:, 1:])\n",
+    "plt.xlabel('iteration')\n",
+    "plt.ylabel('loss')\n",
+    "plt.legend([\"ae\", \"sindy\", \"coef\"])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "82af6aa5",
+   "metadata": {},
+   "source": [
+    "Obtain the predicted solution"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4692510d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tgrid = np.linspace(0, (nt-1)*dt, nt)\n",
+    "\n",
+    "d_Z = d_ae.encoder(d_Xtrain)\n",
+    "Z = d_Z.cpu().detach().numpy()\n",
+    "\n",
+    "Zpred = np.zeros([n_train, physics.nt, ae.n_z])\n",
+    "for case_idx in range(len(coefs)):\n",
+    "    Zpred[case_idx] = ld.simulate(coefs[case_idx], Z[case_idx, 0], tgrid)\n",
+    "\n",
+    "d_Zpred = torch.Tensor(Zpred).to(device)\n",
+    "X_pred = d_ae.decoder(d_Zpred).cpu().detach().numpy()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "38cafdda",
+   "metadata": {},
+   "source": [
+    "Plot the latent variable evolution and compare the prediction with the truth"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3f178a53",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "colors = ['r','g','b','c','m']\n",
+    "\n",
+    "for case_idx in range(len(coefs)):\n",
+    "    plt.figure(1+case_idx, figsize=(8,6))\n",
+    "    for k, color in enumerate(colors):\n",
+    "        plt.plot(Zpred[case_idx][:, k],'-'+color)\n",
+    "        plt.plot(Z[case_idx][:, k],'--'+color)\n",
+    "    plt.ylabel('$Z$')\n",
+    "    plt.xlabel('$t$ ($\\\\times\\Delta t$)')\n",
+    "    plt.legend(['pred','truth'])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cfb4c0bf",
+   "metadata": {},
+   "source": [
+    "Visualization of full-order model solution.\n",
+    "Feel free to change as you see fit."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "da08d933",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%matplotlib inline\n",
+    "\n",
+    "case_idx = 0\n",
+    "time_idx = 14\n",
+    "var_idx = 0\n",
+    "\n",
+    "truth = snapshots[case_idx, time_idx, var_idx]\n",
+    "pred = X_pred[case_idx, time_idx, var_idx]\n",
+    "\n",
+    "plt.figure(1)\n",
+    "plt.contourf(physics.x_grid[0], physics.x_grid[1], truth, 200)\n",
+    "\n",
+    "plt.figure(2)\n",
+    "plt.contourf(physics.x_grid[0], physics.x_grid[1], pred, 200)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "71653612",
+   "metadata": {},
+   "source": [
+    "## Training: (2) Using built-in GPLaSDI trainer\n",
+    "\n",
+    "**Work in progress**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2ad84e16",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from lasdi.gplasdi import BayesianGLaSDI\n",
+    "\n",
+    "trainer_options = {'n_samples': 20,         # number of samples when choosing a new parameter point. Not used in this case\n",
+    "                   'lr': lr,                # learning rate\n",
+    "                   'n_iter': 10000,         # number of iteration in training\n",
+    "                   'max_greedy_iter': 0,    # maximum iteration to perform greedy sampling. Not used in this case\n",
+    "                   'n_greedy': -1,          # frequency of greedy sampling. Not used in this case\n",
+    "                   'sindy_weight': sindy_weight, # weight for latent dynamics loss\n",
+    "                   'coef_weight': coef_weight,   # weight for latent dynamics coefficient regularization\n",
+    "                   'device': 'cpu',              # device to perform training (cpu, cuda, mps)\n",
+    "                   'path_checkpoint': 'checkpoints', # directory to save training checkpoint files\n",
+    "                   'path_results': 'results',        # directory to save training results\n",
+    "                   }\n",
+    "\n",
+    "trainer = BayesianGLaSDI(physics, ae, ld, trainer_options)\n",
+    "trainer.X_train = X_train\n",
+    "\n",
+    "trainer.train()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv_gpu",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}