train-epns.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a618e541-0ac0-41dc-8bbb-a64299a61fdf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/h5/syhu386f/EPNS\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/software/util/JupyterLab/alpha/share/pytorch_v2/lib/python3.10/site-packages/IPython/core/magics/osm.py:417: UserWarning: using dhist requires you to install the `pickleshare` library.\n",
      "  self.shell.db['dhist'] = compress_dhist(dhist)[-100:]\n"
     ]
    }
   ],
   "source": [
    "%cd EPNS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4a1e1d9a-e23e-44d4-a3cd-48b909765953",
   "metadata": {},
   "outputs": [],
   "source": [
    "!source .venv/bin/activate"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7544d447-916e-4a2c-8652-0a73eb02f898",
   "metadata": {},
   "source": [
    "# EPNS Cellular Dynamics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7145fe3-3780-4c63-805f-b6ad38482923",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python train_model.py cell_dynamics_EPNS \\\n",
    "    --data_directory=../simulations/cellular_dynamics \\\n",
    "    --save_path=../models/cellular-dynamics.pt \\\n",
    "    --device=cuda"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7eb7a406-974e-42ee-8864-296fd10dbd07",
   "metadata": {},
   "source": [
    "# Morpheus"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74ead912-5eaf-468c-ba7a-405af27e29b5",
   "metadata": {},
   "source": [
    "## Single Cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24ec5a89-5fec-446d-95dc-fcd31f70e4e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "wandb not found!\n",
      "set data_directory to ../simulations/morpheus-singlecell in main config!\n",
      "set save_path to ../models/morpheus-singlecell-800.pt in main config!\n",
      "set device to cuda in main config!\n",
      "starting new model training run at 1718230443.9956343\n",
      "shape of first datapoint: (60, 80, 80, 2) . Number of train datapoints: 800 number of val datapoints: 100 number of test datapoints: 100\n",
      "/software/util/JupyterLab/alpha/share/pytorch_v2/lib/python3.10/site-packages/torch/nn/modules/lazy.py:181: UserWarning: Lazy modules are a new feature under heavy development so changes to the API or functionality can happen at any moment.\n",
      "  warnings.warn('Lazy modules are a new feature under heavy development '\n",
      "using conditional VAE decoder\n",
      "lowest level shape in U-net processor: torch.Size([3, 256, 10, 10])\n",
      "griddeepset permutation equivariance check succesful!\n",
      "lowest level shape in U-net processor: torch.Size([3, 256, 10, 10])\n",
      "griddeepset permutation equivariance check succesful!\n",
      "the model has 13274689 parameters.\n",
      "will train for 180 epochs with multi-step training.\n",
      "will save model state dict as ../models/morpheus-singlecell-800.pt\n",
      "using multi-step training!\n",
      "training rollout len is 1. Beta: 2.8444444444444446\n",
      "epoch 0 \t loss 7469.390081 \t KL 30.5556866979599 \t mean correct 1.389755\t val loss 1168.59458984375 \t val KL 5.282953262329102 \t val mean correct 1.931106\t took 58.393482 seconds\n",
      "training rollout len is 1. Beta: 5.688888888888889\n",
      "epoch 1 \t loss 850.276959 \t KL 3.001676568984985 \t mean correct 1.948448\t val loss 628.0929248046875 \t val KL 3.1185896396636963 \t val mean correct 1.962256\t took 41.777758 seconds\n",
      "training rollout len is 1. Beta: 8.533333333333333\n",
      "epoch 2 \t loss 544.216201 \t KL 2.9873141956329348 \t mean correct 1.970764\t val loss 408.82041259765623 \t val KL 3.520761013031006 \t val mean correct 1.976572\t took 40.165299 seconds\n",
      "training rollout len is 1. Beta: 11.377777777777778\n",
      "epoch 3 \t loss 399.095586 \t KL 3.2143046641349793 \t mean correct 1.979931\t val loss 319.9370910644531 \t val KL 2.873833656311035 \t val mean correct 1.982378\t took 42.319217 seconds\n",
      "training rollout len is 1. Beta: 14.222222222222221\n",
      "epoch 4 \t loss 341.598951 \t KL 2.8856445574760436 \t mean correct 1.983120\t val loss 294.75732421875 \t val KL 2.978806495666504 \t val mean correct 1.982806\t took 38.427398 seconds\n",
      "training rollout len is 1. Beta: 17.066666666666666\n",
      "epoch 5 \t loss 319.769192 \t KL 2.971522214412689 \t mean correct 1.984429\t val loss 263.6446295166016 \t val KL 3.520192861557007 \t val mean correct 1.984569\t took 44.413214 seconds\n",
      "training rollout len is 1. Beta: 19.91111111111111\n",
      "epoch 6 \t loss 294.965806 \t KL 2.5408189415931703 \t mean correct 1.985150\t val loss 230.003583984375 \t val KL 2.2654991149902344 \t val mean correct 1.985844\t took 47.766156 seconds\n",
      "training rollout len is 1. Beta: 22.755555555555556\n",
      "epoch 7 \t loss 276.835184 \t KL 2.268388465642929 \t mean correct 1.985802\t val loss 225.99728088378907 \t val KL 1.9374418258666992 \t val mean correct 1.986112\t took 39.779138 seconds\n",
      "training rollout len is 1. Beta: 25.6\n",
      "epoch 8 \t loss 256.327684 \t KL 1.6485049510002137 \t mean correct 1.986220\t val loss 211.02795288085937 \t val KL 2.3453407287597656 \t val mean correct 1.986709\t took 38.388979 seconds\n",
      "training rollout len is 1. Beta: 28.444444444444443\n",
      "epoch 9 \t loss 273.036281 \t KL 2.2898140943050382 \t mean correct 1.986447\t val loss 188.0410577392578 \t val KL 2.434065580368042 \t val mean correct 1.987684\t took 43.052293 seconds\n",
      "training rollout len is 1. Beta: 31.288888888888888\n",
      "epoch 10 \t loss 256.210818 \t KL 1.6087611389160157 \t mean correct 1.986480\t val loss 184.52422485351562 \t val KL 1.4197208881378174 \t val mean correct 1.988206\t took 42.480911 seconds\n",
      "training rollout len is 1. Beta: 34.13333333333333\n",
      "epoch 11 \t loss 268.939881 \t KL 1.936484581232071 \t mean correct 1.986605\t val loss 203.79295654296874 \t val KL 1.7633814811706543 \t val mean correct 1.986900\t took 40.849096 seconds\n",
      "training rollout len is 1. Beta: 36.977777777777774\n",
      "epoch 12 \t loss 278.972755 \t KL 2.1503100526332855 \t mean correct 1.986846\t val loss 203.39057250976563 \t val KL 2.486168622970581 \t val mean correct 1.986584\t took 41.636810 seconds\n",
      "training rollout len is 1. Beta: 39.82222222222222\n",
      "epoch 13 \t loss 260.891510 \t KL 1.6384689795970917 \t mean correct 1.986979\t val loss 182.093125 \t val KL 1.710686445236206 \t val mean correct 1.987806\t took 40.431802 seconds\n",
      "training rollout len is 1. Beta: 42.666666666666664\n",
      "epoch 14 \t loss 266.978847 \t KL 1.8406717312335967 \t mean correct 1.987411\t val loss 172.21936462402343 \t val KL 2.2169597148895264 \t val mean correct 1.988544\t took 42.941507 seconds\n",
      "training rollout len is 1. Beta: 45.51111111111111\n",
      "epoch 15 \t loss 254.661443 \t KL 1.4628567206859588 \t mean correct 1.987380\t val loss 182.54972106933593 \t val KL 1.3995039463043213 \t val mean correct 1.987875\t took 41.179756 seconds\n",
      "training rollout len is 1. Beta: 48.355555555555554\n",
      "epoch 16 \t loss 282.044649 \t KL 1.9377081096172333 \t mean correct 1.987348\t val loss 175.1264727783203 \t val KL 2.420921802520752 \t val mean correct 1.988287\t took 48.245517 seconds\n",
      "training rollout len is 1. Beta: 51.2\n",
      "epoch 17 \t loss 283.999226 \t KL 1.8657562553882598 \t mean correct 1.987307\t val loss 179.99230102539062 \t val KL 2.413905620574951 \t val mean correct 1.988019\t took 43.660114 seconds\n",
      "training rollout len is 1. Beta: 54.044444444444444\n",
      "epoch 18 \t loss 265.994946 \t KL 1.6004616057872771 \t mean correct 1.987905\t val loss 168.3944287109375 \t val KL 1.62283456325531 \t val mean correct 1.988856\t took 39.273182 seconds\n",
      "training rollout len is 1. Beta: 56.888888888888886\n",
      "epoch 19 \t loss 287.367080 \t KL 1.8783427333831788 \t mean correct 1.987743\t val loss 177.40212036132812 \t val KL 1.4934219121932983 \t val mean correct 1.988125\t took 37.867855 seconds\n",
      "training rollout len is 1. Beta: 59.733333333333334\n",
      "epoch 20 \t loss 289.328041 \t KL 1.8131067967414856 \t mean correct 1.987713\t val loss 162.2765771484375 \t val KL 1.8295197486877441 \t val mean correct 1.989441\t took 40.293777 seconds\n",
      "training rollout len is 1. Beta: 62.577777777777776\n",
      "epoch 21 \t loss 287.782819 \t KL 1.8343985438346864 \t mean correct 1.988246\t val loss 167.8855059814453 \t val KL 1.8712785243988037 \t val mean correct 1.988737\t took 42.465925 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 22 \t loss 307.114100 \t KL 2.0438950419425965 \t mean correct 1.987998\t val loss 169.2580725097656 \t val KL 2.2841286659240723 \t val mean correct 1.988847\t took 39.648655 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 23 \t loss 316.793313 \t KL 2.1196596240997314 \t mean correct 1.987628\t val loss 170.46387390136718 \t val KL 1.6422762870788574 \t val mean correct 1.988509\t took 52.539970 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 24 \t loss 291.383899 \t KL 1.788702529668808 \t mean correct 1.987865\t val loss 174.93286987304688 \t val KL 1.6744780540466309 \t val mean correct 1.988559\t took 40.175225 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 25 \t loss 336.720581 \t KL 2.253477039337158 \t mean correct 1.987569\t val loss 578.587373046875 \t val KL 1.8238048553466797 \t val mean correct 1.980569\t took 40.669729 seconds\n",
      "val rollout loss: nan, val rollout mean correct: 1.883719429805361, took 77.823309 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 26 \t loss 1499.108129 \t KL 1.5884530866146087 \t mean correct 1.979608\t val loss 361.4845068359375 \t val KL 2.0477294921875 \t val mean correct 1.981991\t took 41.502028 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 27 \t loss 718.327239 \t KL 2.0188747000694276 \t mean correct 1.984254\t val loss 217.04004028320313 \t val KL 1.2925459146499634 \t val mean correct 1.986403\t took 42.415816 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 28 \t loss 305.957499 \t KL 1.781403774023056 \t mean correct 1.987114\t val loss 171.93180236816406 \t val KL 1.7785383462905884 \t val mean correct 1.988422\t took 45.282793 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 29 \t loss 316.257064 \t KL 2.149212030172348 \t mean correct 1.987816\t val loss 161.276416015625 \t val KL 1.574447751045227 \t val mean correct 1.989222\t took 42.596959 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 30 \t loss 315.659808 \t KL 2.1764649641513825 \t mean correct 1.987973\t val loss 173.160244140625 \t val KL 2.1370861530303955 \t val mean correct 1.988387\t took 42.285361 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 31 \t loss 317.898111 \t KL 2.284200813770294 \t mean correct 1.988303\t val loss 173.2819140625 \t val KL 2.160762071609497 \t val mean correct 1.988572\t took 41.415419 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 32 \t loss 299.836929 \t KL 1.9188957369327546 \t mean correct 1.987858\t val loss 168.22055847167968 \t val KL 2.21891188621521 \t val mean correct 1.989059\t took 39.514138 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 33 \t loss 320.101271 \t KL 2.283901489973068 \t mean correct 1.988107\t val loss 159.02355712890625 \t val KL 2.3687639236450195 \t val mean correct 1.989253\t took 39.958081 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 34 \t loss 316.656198 \t KL 2.1608244836330415 \t mean correct 1.987760\t val loss 175.00944976806642 \t val KL 2.0988996028900146 \t val mean correct 1.988103\t took 40.503147 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 35 \t loss 327.802613 \t KL 2.281126172542572 \t mean correct 1.987642\t val loss 186.64053649902343 \t val KL 1.8847240209579468 \t val mean correct 1.987444\t took 44.585883 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 36 \t loss 342.756354 \t KL 2.514453842639923 \t mean correct 1.987661\t val loss 175.43258666992188 \t val KL 2.5265376567840576 \t val mean correct 1.988478\t took 41.871781 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 37 \t loss 312.291346 \t KL 2.1054905641078947 \t mean correct 1.988028\t val loss 182.7081756591797 \t val KL 2.625011444091797 \t val mean correct 1.987781\t took 39.379101 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 38 \t loss 352.367764 \t KL 2.571348559856415 \t mean correct 1.987358\t val loss 204.63052307128908 \t val KL 2.211688995361328 \t val mean correct 1.986697\t took 44.262051 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 39 \t loss 319.401057 \t KL 2.0817503833770754 \t mean correct 1.987475\t val loss 179.7070733642578 \t val KL 1.6060678958892822 \t val mean correct 1.988016\t took 40.284279 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 40 \t loss 315.419871 \t KL 2.0334406077861784 \t mean correct 1.987527\t val loss 168.7519372558594 \t val KL 1.8985536098480225 \t val mean correct 1.988700\t took 45.662994 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 41 \t loss 317.410341 \t KL 2.1249904572963714 \t mean correct 1.987742\t val loss 167.25881774902345 \t val KL 2.6992807388305664 \t val mean correct 1.988875\t took 41.891287 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 42 \t loss 300.408060 \t KL 1.9684361743927001 \t mean correct 1.988155\t val loss 170.4458349609375 \t val KL 1.5179226398468018 \t val mean correct 1.988506\t took 40.699190 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 43 \t loss 310.849686 \t KL 2.231839416027069 \t mean correct 1.988548\t val loss 172.779150390625 \t val KL 2.9193713665008545 \t val mean correct 1.988578\t took 41.073850 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 44 \t loss 320.957187 \t KL 2.317637702226639 \t mean correct 1.988276\t val loss 157.739931640625 \t val KL 2.2608461380004883 \t val mean correct 1.989597\t took 39.454508 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 45 \t loss 332.222685 \t KL 2.479104653596878 \t mean correct 1.988157\t val loss 171.12192016601563 \t val KL 2.1437857151031494 \t val mean correct 1.988856\t took 40.435400 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 46 \t loss 356.850379 \t KL 2.534430285692215 \t mean correct 1.987385\t val loss 180.8109164428711 \t val KL 1.7471325397491455 \t val mean correct 1.988194\t took 45.653255 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 47 \t loss 338.686401 \t KL 2.0667143666744234 \t mean correct 1.986853\t val loss 248.3595751953125 \t val KL 2.066702127456665 \t val mean correct 1.986019\t took 50.325422 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 48 \t loss 302.394307 \t KL 1.8741175734996796 \t mean correct 1.987925\t val loss 167.19604309082033 \t val KL 2.1772844791412354 \t val mean correct 1.988878\t took 44.875950 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 49 \t loss 306.069814 \t KL 2.0980682408809663 \t mean correct 1.988381\t val loss 192.5044677734375 \t val KL 1.8166848421096802 \t val mean correct 1.987016\t took 41.408853 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 50 \t loss 295.473873 \t KL 1.855777782201767 \t mean correct 1.988100\t val loss 174.1023486328125 \t val KL 2.0683696269989014 \t val mean correct 1.988125\t took 43.230661 seconds\n",
      "val rollout loss: nan, val rollout mean correct: 1.9036670342807112, took 49.725612 seconds\n",
      "training rollout len is 2. Beta: 64.0\n",
      "epoch 51 \t loss 347.185503 \t KL 1.791843365430832 \t mean correct 1.984652\t val loss 183.90722106933595 \t val KL 1.5817257165908813 \t val mean correct 1.988797\t took 55.215005 seconds\n",
      "training rollout len is 2. Beta: 64.0\n",
      "epoch 52 \t loss 326.784111 \t KL 1.602620245218277 \t mean correct 1.985240\t val loss 161.11854125976564 \t val KL 1.6680207252502441 \t val mean correct 1.989469\t took 55.788171 seconds\n",
      "training rollout len is 2. Beta: 64.0\n",
      "epoch 53 \t loss 336.079543 \t KL 1.723378536105156 \t mean correct 1.985167\t val loss 164.70963806152344 \t val KL 1.8195348978042603 \t val mean correct 1.989797\t took 54.702402 seconds\n",
      "training rollout len is 2. Beta: 64.0\n",
      "epoch 54 \t loss 348.177788 \t KL 1.928271571993828 \t mean correct 1.985074\t val loss 173.2772088623047 \t val KL 2.0711593627929688 \t val mean correct 1.988925\t took 51.445480 seconds\n",
      "training rollout len is 2. Beta: 64.0\n",
      "epoch 55 \t loss 362.898104 \t KL 2.2865602374076843 \t mean correct 1.985646\t val loss 181.5552313232422 \t val KL 2.435051441192627 \t val mean correct 1.988444\t took 48.440409 seconds\n",
      "training rollout len is 2. Beta: 64.0\n",
      "epoch 56 \t loss 386.130964 \t KL 2.4405964756011964 \t mean correct 1.984755\t val loss 181.8715155029297 \t val KL 2.445345640182495 \t val mean correct 1.988228\t took 49.725583 seconds\n",
      "training rollout len is 2. Beta: 64.0\n",
      "epoch 57 \t loss 384.610415 \t KL 2.568719345331192 \t mean correct 1.985410\t val loss 178.59840942382812 \t val KL 2.602588176727295 \t val mean correct 1.988322\t took 51.185481 seconds\n",
      "training rollout len is 2. Beta: 64.0\n",
      "epoch 58 \t loss 389.017887 \t KL 2.6170005989074707 \t mean correct 1.985409\t val loss 177.26915405273436 \t val KL 2.7271275520324707 \t val mean correct 1.988628\t took 48.392602 seconds\n",
      "training rollout len is 2. Beta: 64.0\n",
      "epoch 59 \t loss 403.430264 \t KL 2.8722948813438416 \t mean correct 1.985564\t val loss 171.5037384033203 \t val KL 2.8607356548309326 \t val mean correct 1.988969\t took 48.938920 seconds\n",
      "training rollout len is 3. Beta: 64.0\n",
      "epoch 60 \t loss 429.597540 \t KL 2.5606332389513655 \t mean correct 1.982689\t val loss 175.034033203125 \t val KL 2.554079532623291 \t val mean correct 1.989297\t took 57.990962 seconds\n",
      "training rollout len is 3. Beta: 64.0\n",
      "epoch 61 \t loss 441.347078 \t KL 2.6042401631673173 \t mean correct 1.982032\t val loss 181.99371826171875 \t val KL 2.386627674102783 \t val mean correct 1.989475\t took 57.623832 seconds\n",
      "training rollout len is 3. Beta: 64.0\n",
      "epoch 62 \t loss 428.174178 \t KL 2.451198116938274 \t mean correct 1.982376\t val loss 180.55025390625 \t val KL 2.51436448097229 \t val mean correct 1.988800\t took 58.306543 seconds\n",
      "training rollout len is 3. Beta: 64.0\n",
      "epoch 63 \t loss 457.895900 \t KL 2.953890020847321 \t mean correct 1.982291\t val loss 206.85793762207032 \t val KL 2.7804720401763916 \t val mean correct 1.989084\t took 59.858999 seconds\n",
      "training rollout len is 3. Beta: 64.0\n",
      "epoch 64 \t loss 448.942053 \t KL 2.7611745564142858 \t mean correct 1.982312\t val loss 174.3084289550781 \t val KL 2.718930721282959 \t val mean correct 1.989544\t took 60.632324 seconds\n",
      "training rollout len is 3. Beta: 64.0\n",
      "epoch 65 \t loss 451.104352 \t KL 2.805943342844645 \t mean correct 1.982317\t val loss 176.40394836425781 \t val KL 3.2981069087982178 \t val mean correct 1.989803\t took 58.407302 seconds\n",
      "training rollout len is 3. Beta: 64.0\n",
      "epoch 66 \t loss 434.865458 \t KL 2.670133345127106 \t mean correct 1.982769\t val loss 185.324794921875 \t val KL 2.4203686714172363 \t val mean correct 1.989403\t took 57.714933 seconds\n",
      "training rollout len is 3. Beta: 64.0\n",
      "epoch 67 \t loss 450.415478 \t KL 2.849340230623882 \t mean correct 1.982346\t val loss 189.009521484375 \t val KL 2.5427255630493164 \t val mean correct 1.989031\t took 57.990411 seconds\n",
      "training rollout len is 3. Beta: 64.0\n",
      "epoch 68 \t loss 425.705146 \t KL 2.3977947449684134 \t mean correct 1.982323\t val loss 205.48881469726564 \t val KL 2.491835832595825 \t val mean correct 1.988953\t took 57.676368 seconds\n",
      "training rollout len is 3. Beta: 64.0\n",
      "epoch 69 \t loss 460.683668 \t KL 2.9340141304334004 \t mean correct 1.982159\t val loss 206.77704223632813 \t val KL 2.745467185974121 \t val mean correct 1.988206\t took 59.492936 seconds\n",
      "training rollout len is 4. Beta: 64.0\n",
      "epoch 70 \t loss 485.811274 \t KL 2.4630586963891985 \t mean correct 1.978841\t val loss 205.89942016601563 \t val KL 2.391727924346924 \t val mean correct 1.987662\t took 67.335092 seconds\n",
      "training rollout len is 4. Beta: 64.0\n",
      "epoch 71 \t loss 485.404661 \t KL 2.515977550148964 \t mean correct 1.979111\t val loss 227.39287475585937 \t val KL 2.9505069255828857 \t val mean correct 1.987206\t took 66.252847 seconds\n",
      "training rollout len is 4. Beta: 64.0\n",
      "epoch 72 \t loss 503.046239 \t KL 2.5769464069604875 \t mean correct 1.978456\t val loss 225.12873474121093 \t val KL 2.4146432876586914 \t val mean correct 1.988572\t took 67.148155 seconds\n",
      "training rollout len is 4. Beta: 64.0\n",
      "epoch 73 \t loss 503.881238 \t KL 2.676163896918297 \t mean correct 1.978522\t val loss 212.3926806640625 \t val KL 2.628537178039551 \t val mean correct 1.988634\t took 68.458239 seconds\n",
      "training rollout len is 4. Beta: 64.0\n",
      "epoch 74 \t loss 508.043175 \t KL 2.64295639872551 \t mean correct 1.978130\t val loss 231.3429217529297 \t val KL 2.778716564178467 \t val mean correct 1.988803\t took 68.235023 seconds\n",
      "training rollout len is 4. Beta: 64.0\n",
      "epoch 75 \t loss 509.680756 \t KL 2.807936898469925 \t mean correct 1.978885\t val loss 221.06937622070313 \t val KL 2.636291980743408 \t val mean correct 1.988775\t took 68.897415 seconds\n",
      "val rollout loss: nan, val rollout mean correct: 1.9057539893840922, took 54.537598 seconds\n",
      "training rollout len is 4. Beta: 64.0\n",
      "epoch 76 \t loss 515.014601 \t KL 2.918522120118141 \t mean correct 1.978872\t val loss 201.8068054199219 \t val KL 2.604367971420288 \t val mean correct 1.988825\t took 74.616909 seconds\n",
      "training rollout len is 4. Beta: 64.0\n",
      "epoch 77 \t loss 495.897967 \t KL 2.7223248994350433 \t mean correct 1.979166\t val loss 198.40421997070314 \t val KL 2.7422995567321777 \t val mean correct 1.988512\t took 67.796823 seconds\n",
      "training rollout len is 4. Beta: 64.0\n",
      "epoch 78 \t loss 492.795553 \t KL 2.519884389042854 \t mean correct 1.978729\t val loss 196.28026306152344 \t val KL 2.442488431930542 \t val mean correct 1.988241\t took 70.062466 seconds\n",
      "training rollout len is 4. Beta: 64.0\n",
      "epoch 79 \t loss 476.030271 \t KL 2.5038559460639953 \t mean correct 1.979465\t val loss 207.13392333984376 \t val KL 2.6126511096954346 \t val mean correct 1.987737\t took 67.198706 seconds\n",
      "training rollout len is 5. Beta: 64.0\n",
      "epoch 80 \t loss 511.833401 \t KL 2.4053288869857785 \t mean correct 1.976995\t val loss 199.86486938476563 \t val KL 2.4742376804351807 \t val mean correct 1.988797\t took 75.968666 seconds\n",
      "training rollout len is 5. Beta: 64.0\n",
      "epoch 81 \t loss 561.445149 \t KL 2.8132332859039315 \t mean correct 1.975796\t val loss 226.13140441894532 \t val KL 2.9371860027313232 \t val mean correct 1.987284\t took 77.668936 seconds\n",
      "training rollout len is 5. Beta: 64.0\n",
      "epoch 82 \t loss 559.670772 \t KL 2.8157708783149724 \t mean correct 1.975838\t val loss 228.27520935058592 \t val KL 2.819159507751465 \t val mean correct 1.988397\t took 79.010383 seconds\n",
      "training rollout len is 5. Beta: 64.0\n",
      "epoch 83 \t loss 587.741833 \t KL 3.1204626045227055 \t mean correct 1.975092\t val loss 215.5458135986328 \t val KL 3.3098678588867188 \t val mean correct 1.987816\t took 76.715413 seconds\n",
      "training rollout len is 5. Beta: 64.0\n",
      "epoch 84 \t loss 558.955941 \t KL 2.9552280521392826 \t mean correct 1.976203\t val loss 195.3374920654297 \t val KL 2.940808057785034 \t val mean correct 1.988725\t took 76.191769 seconds\n",
      "training rollout len is 5. Beta: 64.0\n",
      "epoch 85 \t loss 600.080933 \t KL 3.1927719826698304 \t mean correct 1.975070\t val loss 218.31422424316406 \t val KL 3.3397741317749023 \t val mean correct 1.986709\t took 78.672879 seconds\n",
      "training rollout len is 5. Beta: 64.0\n",
      "epoch 86 \t loss 1106.623242 \t KL 3.38718738746643 \t mean correct 1.962051\t val loss 3577.78671875 \t val KL 3.277118444442749 \t val mean correct 1.900931\t took 76.675246 seconds\n",
      "training rollout len is 5. Beta: 64.0\n",
      "epoch 87 \t loss 1833.564520 \t KL 2.7257377963066096 \t mean correct 1.945088\t val loss 208.08033264160156 \t val KL 2.377488613128662 \t val mean correct 1.987869\t took 76.577561 seconds\n",
      "training rollout len is 5. Beta: 64.0\n",
      "epoch 88 \t loss 539.226740 \t KL 2.397045279979706 \t mean correct 1.975789\t val loss 243.49507263183594 \t val KL 2.7379982471466064 \t val mean correct 1.987397\t took 80.407589 seconds\n",
      "training rollout len is 5. Beta: 64.0\n",
      "epoch 89 \t loss 570.763321 \t KL 2.8960123109817517 \t mean correct 1.975862\t val loss 197.53852600097656 \t val KL 2.6442723274230957 \t val mean correct 1.987853\t took 78.221230 seconds\n",
      "training rollout len is 6. Beta: 64.0\n",
      "epoch 90 \t loss 601.533851 \t KL 2.509352062940597 \t mean correct 1.972361\t val loss 275.46587036132814 \t val KL 2.6141562461853027 \t val mean correct 1.985950\t took 86.348860 seconds\n",
      "training rollout len is 6. Beta: 64.0\n",
      "epoch 91 \t loss 657.922584 \t KL 2.8131655557950337 \t mean correct 1.970751\t val loss 253.50049987792968 \t val KL 2.5324018001556396 \t val mean correct 1.986537\t took 86.474757 seconds\n",
      "training rollout len is 6. Beta: 64.0\n",
      "epoch 92 \t loss 618.448424 \t KL 2.6141936083634687 \t mean correct 1.971662\t val loss 258.8987182617187 \t val KL 2.6620399951934814 \t val mean correct 1.987953\t took 87.943951 seconds\n",
      "training rollout len is 6. Beta: 64.0\n",
      "epoch 93 \t loss 654.596457 \t KL 3.0879087467988326 \t mean correct 1.971331\t val loss 227.32601196289062 \t val KL 3.037238121032715 \t val mean correct 1.987091\t took 85.757749 seconds\n",
      "training rollout len is 6. Beta: 64.0\n",
      "epoch 94 \t loss 650.805906 \t KL 3.0501141115029653 \t mean correct 1.971593\t val loss 263.64301147460935 \t val KL 3.1985814571380615 \t val mean correct 1.985731\t took 86.857287 seconds\n",
      "training rollout len is 6. Beta: 64.0\n",
      "epoch 95 \t loss 611.212747 \t KL 2.8071595780054723 \t mean correct 1.972913\t val loss 272.2175146484375 \t val KL 2.773655891418457 \t val mean correct 1.987091\t took 86.654094 seconds\n",
      "training rollout len is 6. Beta: 64.0\n",
      "epoch 96 \t loss 632.220531 \t KL 2.852935330470403 \t mean correct 1.972427\t val loss 272.4062609863281 \t val KL 2.815610885620117 \t val mean correct 1.985400\t took 87.187304 seconds\n",
      "training rollout len is 6. Beta: 64.0\n",
      "epoch 97 \t loss 647.953633 \t KL 2.9229659704367315 \t mean correct 1.971185\t val loss 240.81945068359374 \t val KL 2.4583590030670166 \t val mean correct 1.987962\t took 87.697801 seconds\n",
      "training rollout len is 6. Beta: 64.0\n",
      "epoch 98 \t loss 590.631494 \t KL 2.552549905379613 \t mean correct 1.972684\t val loss 318.34124633789065 \t val KL 2.9036245346069336 \t val mean correct 1.985078\t took 86.909302 seconds\n",
      "training rollout len is 6. Beta: 64.0\n",
      "epoch 99 \t loss 610.307319 \t KL 2.92265766620636 \t mean correct 1.973453\t val loss 216.6892333984375 \t val KL 3.2263381481170654 \t val mean correct 1.987991\t took 88.363009 seconds\n",
      "training rollout len is 7. Beta: 64.0\n",
      "epoch 100 \t loss 662.974016 \t KL 3.0020772743225104 \t mean correct 1.970469\t val loss 329.7746154785156 \t val KL 2.7721011638641357 \t val mean correct 1.987025\t took 96.927054 seconds\n",
      "val rollout loss: nan, val rollout mean correct: 1.899521405450229, took 52.991954 seconds\n",
      "training rollout len is 7. Beta: 64.0\n",
      "epoch 101 \t loss 661.474825 \t KL 2.991099109990257 \t mean correct 1.970474\t val loss 236.16673400878906 \t val KL 2.974499464035034 \t val mean correct 1.987797\t took 97.622692 seconds\n",
      "training rollout len is 7. Beta: 64.0\n",
      "epoch 102 \t loss 664.394526 \t KL 2.954290104934148 \t mean correct 1.969898\t val loss 264.3866888427734 \t val KL 3.173459053039551 \t val mean correct 1.987612\t took 97.823668 seconds\n",
      "training rollout len is 7. Beta: 64.0\n",
      "epoch 103 \t loss 682.096406 \t KL 3.2913285596030097 \t mean correct 1.970057\t val loss 259.6706262207031 \t val KL 3.165576934814453 \t val mean correct 1.986697\t took 96.776758 seconds\n",
      "training rollout len is 7. Beta: 64.0\n",
      "epoch 104 \t loss 699.685845 \t KL 3.1884827474185404 \t mean correct 1.969306\t val loss 268.6662744140625 \t val KL 3.1885180473327637 \t val mean correct 1.988103\t took 97.909032 seconds\n",
      "training rollout len is 7. Beta: 64.0\n"
     ]
    }
   ],
   "source": [
    "!python train_model.py cell_dynamics_EPNS \\\n",
    "    --data_directory=../simulations/morpheus-singlecell \\\n",
    "    --save_path=../models/morpheus-singlecell-800.pt \\\n",
    "    --device=cuda"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "553948d6-a17d-4ea6-bc68-8eccf34ed461",
   "metadata": {},
   "source": [
    "## Two Cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4163e463-556c-4f9f-bb7f-2acb2b9424eb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "wandb not found!\n",
      "set data_directory to ../simulations/morpheus-twocell in main config!\n",
      "set save_path to ../models/morpheus-twocell-800.pt in main config!\n",
      "set device to cuda in main config!\n",
      "starting new model training run at 1718257131.4292805\n",
      "shape of first datapoint: (60, 80, 80, 2) . Number of train datapoints: 800 number of val datapoints: 100 number of test datapoints: 100\n",
      "/software/util/JupyterLab/alpha/share/pytorch_v2/lib/python3.10/site-packages/torch/nn/modules/lazy.py:181: UserWarning: Lazy modules are a new feature under heavy development so changes to the API or functionality can happen at any moment.\n",
      "  warnings.warn('Lazy modules are a new feature under heavy development '\n",
      "using conditional VAE decoder\n",
      "lowest level shape in U-net processor: torch.Size([4, 256, 10, 10])\n",
      "griddeepset permutation equivariance check succesful!\n",
      "lowest level shape in U-net processor: torch.Size([4, 256, 10, 10])\n",
      "griddeepset permutation equivariance check succesful!\n",
      "the model has 13274689 parameters.\n",
      "will train for 180 epochs with multi-step training.\n",
      "will save model state dict as ../models/morpheus-twocell-800.pt\n",
      "using multi-step training!\n",
      "training rollout len is 1. Beta: 2.8444444444444446\n",
      "epoch 0 \t loss 8184.818112 \t KL 23.73462411880493 \t mean correct 1.444279\t val loss 1263.0315478515624 \t val KL 6.297850131988525 \t val mean correct 1.925548\t took 33.391516 seconds\n",
      "training rollout len is 1. Beta: 5.688888888888889\n",
      "epoch 1 \t loss 1103.357314 \t KL 3.4720298218727113 \t mean correct 1.933386\t val loss 889.0200366210937 \t val KL 2.906303882598877 \t val mean correct 1.945809\t took 12.073235 seconds\n",
      "training rollout len is 1. Beta: 8.533333333333333\n",
      "epoch 2 \t loss 817.818776 \t KL 3.380333893299103 \t mean correct 1.953593\t val loss 685.7878076171875 \t val KL 3.430026054382324 \t val mean correct 1.959844\t took 12.235091 seconds\n",
      "training rollout len is 1. Beta: 11.377777777777778\n",
      "epoch 3 \t loss 642.114575 \t KL 3.7433365964889527 \t mean correct 1.963820\t val loss 525.2360559082031 \t val KL 4.339087963104248 \t val mean correct 1.968348\t took 12.160542 seconds\n",
      "training rollout len is 1. Beta: 14.222222222222221\n",
      "epoch 4 \t loss 564.571235 \t KL 3.824009220600128 \t mean correct 1.968652\t val loss 485.89897094726564 \t val KL 3.5755105018615723 \t val mean correct 1.969419\t took 11.936914 seconds\n",
      "training rollout len is 1. Beta: 17.066666666666666\n",
      "epoch 5 \t loss 530.389382 \t KL 3.6882430148124694 \t mean correct 1.970588\t val loss 458.8078564453125 \t val KL 4.504819393157959 \t val mean correct 1.970914\t took 12.022219 seconds\n",
      "training rollout len is 1. Beta: 19.91111111111111\n",
      "epoch 6 \t loss 514.473949 \t KL 3.422036271095276 \t mean correct 1.971400\t val loss 437.0382238769531 \t val KL 3.72861909866333 \t val mean correct 1.972245\t took 11.935523 seconds\n",
      "training rollout len is 1. Beta: 22.755555555555556\n",
      "epoch 7 \t loss 503.486691 \t KL 3.3420767259597777 \t mean correct 1.972170\t val loss 421.22509765625 \t val KL 3.074723958969116 \t val mean correct 1.972031\t took 12.034111 seconds\n",
      "training rollout len is 1. Beta: 25.6\n",
      "epoch 8 \t loss 489.762025 \t KL 2.971195878982544 \t mean correct 1.973046\t val loss 401.5779089355469 \t val KL 2.022395372390747 \t val mean correct 1.973955\t took 12.075976 seconds\n",
      "training rollout len is 1. Beta: 28.444444444444443\n",
      "epoch 9 \t loss 477.509026 \t KL 2.6747398674488068 \t mean correct 1.973518\t val loss 398.5549169921875 \t val KL 2.323073148727417 \t val mean correct 1.973869\t took 11.951718 seconds\n",
      "training rollout len is 1. Beta: 31.288888888888888\n",
      "epoch 10 \t loss 496.342939 \t KL 3.055432720184326 \t mean correct 1.973545\t val loss 396.8874853515625 \t val KL 2.5575759410858154 \t val mean correct 1.973550\t took 11.958643 seconds\n",
      "training rollout len is 1. Beta: 34.13333333333333\n",
      "epoch 11 \t loss 496.915558 \t KL 3.0733509993553163 \t mean correct 1.973942\t val loss 388.63318603515626 \t val KL 3.4494388103485107 \t val mean correct 1.974142\t took 11.992422 seconds\n",
      "training rollout len is 1. Beta: 36.977777777777774\n",
      "epoch 12 \t loss 509.385691 \t KL 3.0455990409851075 \t mean correct 1.973547\t val loss 387.19028442382813 \t val KL 3.506974458694458 \t val mean correct 1.974434\t took 11.975101 seconds\n",
      "training rollout len is 1. Beta: 39.82222222222222\n",
      "epoch 13 \t loss 492.361604 \t KL 2.5391018736362456 \t mean correct 1.974137\t val loss 395.6889123535156 \t val KL 1.5839756727218628 \t val mean correct 1.973992\t took 11.996294 seconds\n",
      "training rollout len is 1. Beta: 42.666666666666664\n",
      "epoch 14 \t loss 486.725151 \t KL 2.5115809655189514 \t mean correct 1.974664\t val loss 381.1575915527344 \t val KL 2.2092933654785156 \t val mean correct 1.975278\t took 11.939357 seconds\n",
      "training rollout len is 1. Beta: 45.51111111111111\n",
      "epoch 15 \t loss 504.200390 \t KL 2.632641817331314 \t mean correct 1.974249\t val loss 379.8773034667969 \t val KL 2.312361240386963 \t val mean correct 1.974709\t took 12.001179 seconds\n",
      "training rollout len is 1. Beta: 48.355555555555554\n",
      "epoch 16 \t loss 512.705301 \t KL 2.7970974922180174 \t mean correct 1.974790\t val loss 353.69732666015625 \t val KL 2.1205832958221436 \t val mean correct 1.976516\t took 11.892403 seconds\n",
      "training rollout len is 1. Beta: 51.2\n",
      "epoch 17 \t loss 507.324117 \t KL 2.4554484486579895 \t mean correct 1.974527\t val loss 380.6955407714844 \t val KL 2.0860466957092285 \t val mean correct 1.974267\t took 12.034613 seconds\n",
      "training rollout len is 1. Beta: 54.044444444444444\n",
      "epoch 18 \t loss 502.828182 \t KL 2.3963264739513397 \t mean correct 1.975071\t val loss 362.6589208984375 \t val KL 2.2744760513305664 \t val mean correct 1.975956\t took 11.921019 seconds\n",
      "training rollout len is 1. Beta: 56.888888888888886\n",
      "epoch 19 \t loss 511.448109 \t KL 2.3995395982265473 \t mean correct 1.974977\t val loss 375.94024536132815 \t val KL 1.9805872440338135 \t val mean correct 1.975687\t took 11.911639 seconds\n",
      "training rollout len is 1. Beta: 59.733333333333334\n",
      "epoch 20 \t loss 474.095060 \t KL 1.7350926756858827 \t mean correct 1.975208\t val loss 371.110029296875 \t val KL 2.422337532043457 \t val mean correct 1.975462\t took 11.884920 seconds\n",
      "training rollout len is 1. Beta: 62.577777777777776\n",
      "epoch 21 \t loss 479.548848 \t KL 1.8219532561302185 \t mean correct 1.975509\t val loss 352.20634033203123 \t val KL 1.6465225219726562 \t val mean correct 1.976344\t took 12.061525 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 22 \t loss 509.844331 \t KL 2.2837244760990143 \t mean correct 1.975575\t val loss 359.74458740234377 \t val KL 1.9434479475021362 \t val mean correct 1.976012\t took 11.964918 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 23 \t loss 497.655848 \t KL 2.084371408224106 \t mean correct 1.975660\t val loss 367.2546044921875 \t val KL 2.9987874031066895 \t val mean correct 1.975617\t took 12.083946 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 24 \t loss 527.762294 \t KL 2.5141011595726015 \t mean correct 1.975361\t val loss 344.37221435546877 \t val KL 3.1465437412261963 \t val mean correct 1.976775\t took 12.077233 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 25 \t loss 509.210068 \t KL 2.301861538887024 \t mean correct 1.975784\t val loss 353.0991528320312 \t val KL 1.5964574813842773 \t val mean correct 1.976422\t took 12.027114 seconds\n",
      "val rollout loss: nan, val rollout mean correct: 1.8242294627222522, took 25.386903 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 26 \t loss 507.276212 \t KL 2.220617197751999 \t mean correct 1.975608\t val loss 363.04368286132814 \t val KL 1.8409545421600342 \t val mean correct 1.975797\t took 11.930827 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 27 \t loss 531.515919 \t KL 2.683168578147888 \t mean correct 1.975844\t val loss 351.89391357421874 \t val KL 3.292694330215454 \t val mean correct 1.976869\t took 12.056798 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 28 \t loss 515.541855 \t KL 2.399479752779007 \t mean correct 1.975805\t val loss 366.0639270019531 \t val KL 1.4252156019210815 \t val mean correct 1.975302\t took 12.003486 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 29 \t loss 499.309295 \t KL 2.1651342415809633 \t mean correct 1.975750\t val loss 354.9915295410156 \t val KL 2.1848456859588623 \t val mean correct 1.975991\t took 11.995853 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 30 \t loss 530.775691 \t KL 2.7356416821479796 \t mean correct 1.976239\t val loss 363.95676513671873 \t val KL 3.6801750659942627 \t val mean correct 1.975803\t took 12.013560 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 31 \t loss 499.687734 \t KL 2.240926958322525 \t mean correct 1.976006\t val loss 351.527001953125 \t val KL 2.289454698562622 \t val mean correct 1.976498\t took 12.003592 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 32 \t loss 515.915016 \t KL 2.436913408041 \t mean correct 1.975873\t val loss 354.68042236328125 \t val KL 2.447000741958618 \t val mean correct 1.976305\t took 11.940240 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 33 \t loss 516.535204 \t KL 2.532372648715973 \t mean correct 1.976187\t val loss 345.2846960449219 \t val KL 2.5199131965637207 \t val mean correct 1.976694\t took 11.930440 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 34 \t loss 529.589587 \t KL 2.790353448390961 \t mean correct 1.976401\t val loss 342.3583093261719 \t val KL 3.7299418449401855 \t val mean correct 1.977187\t took 11.894725 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 35 \t loss 471.527833 \t KL 1.8687984824180603 \t mean correct 1.976293\t val loss 349.16515686035154 \t val KL 2.4067506790161133 \t val mean correct 1.976944\t took 11.916495 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 36 \t loss 517.979765 \t KL 2.614976415634155 \t mean correct 1.976419\t val loss 351.51653930664065 \t val KL 2.2525525093078613 \t val mean correct 1.976648\t took 11.934652 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 37 \t loss 518.565660 \t KL 2.5098718285560606 \t mean correct 1.975875\t val loss 340.3020300292969 \t val KL 2.132723331451416 \t val mean correct 1.977203\t took 11.890517 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 38 \t loss 529.081835 \t KL 2.7962468528747557 \t mean correct 1.976546\t val loss 338.9749328613281 \t val KL 2.8146073818206787 \t val mean correct 1.977491\t took 11.924775 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 39 \t loss 521.514211 \t KL 2.5259563755989074 \t mean correct 1.975848\t val loss 368.2942419433594 \t val KL 2.567904233932495 \t val mean correct 1.975447\t took 11.921354 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 40 \t loss 527.426321 \t KL 2.7328800439834593 \t mean correct 1.976358\t val loss 352.62760986328124 \t val KL 3.2373876571655273 \t val mean correct 1.976314\t took 12.023335 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 41 \t loss 527.461875 \t KL 2.743197855949402 \t mean correct 1.976425\t val loss 342.083369140625 \t val KL 2.5919835567474365 \t val mean correct 1.977255\t took 11.914328 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 42 \t loss 527.644438 \t KL 2.839701519012451 \t mean correct 1.976797\t val loss 337.806591796875 \t val KL 2.9644224643707275 \t val mean correct 1.977453\t took 12.037997 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 43 \t loss 532.933645 \t KL 2.844277148246765 \t mean correct 1.976546\t val loss 337.10040405273435 \t val KL 2.414952278137207 \t val mean correct 1.977452\t took 11.922429 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 44 \t loss 537.055451 \t KL 2.943623993396759 \t mean correct 1.976551\t val loss 347.9767785644531 \t val KL 3.6857829093933105 \t val mean correct 1.976836\t took 11.919086 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 45 \t loss 524.927116 \t KL 2.8192554783821104 \t mean correct 1.976866\t val loss 358.905732421875 \t val KL 3.3042635917663574 \t val mean correct 1.976267\t took 11.976816 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 46 \t loss 528.129572 \t KL 2.816543686389923 \t mean correct 1.976550\t val loss 340.2524133300781 \t val KL 4.187098026275635 \t val mean correct 1.977255\t took 11.973736 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 47 \t loss 547.184247 \t KL 3.0693617844581604 \t mean correct 1.976482\t val loss 359.51474853515623 \t val KL 2.5576388835906982 \t val mean correct 1.976244\t took 11.925316 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 48 \t loss 511.109288 \t KL 2.6576802587509154 \t mean correct 1.976919\t val loss 357.79454711914065 \t val KL 2.0976510047912598 \t val mean correct 1.976098\t took 11.880543 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 49 \t loss 516.458713 \t KL 2.6358393669128417 \t mean correct 1.976680\t val loss 329.6909033203125 \t val KL 2.2559988498687744 \t val mean correct 1.978002\t took 11.948683 seconds\n",
      "training rollout len is 1. Beta: 64.0\n",
      "epoch 50 \t loss 510.598411 \t KL 2.561050021648407 \t mean correct 1.976628\t val loss 348.63527221679686 \t val KL 3.4247310161590576 \t val mean correct 1.976784\t took 11.951114 seconds\n",
      "val rollout loss: nan, val rollout mean correct: 1.7915463730384564, took 25.257996 seconds\n",
      "training rollout len is 2. Beta: 64.0\n",
      "epoch 51 \t loss 649.755088 \t KL 3.183178687095642 \t mean correct 1.970757\t val loss 363.31303955078124 \t val KL 2.5715246200561523 \t val mean correct 1.976000\t took 19.615223 seconds\n",
      "training rollout len is 2. Beta: 64.0\n",
      "epoch 52 \t loss 649.831325 \t KL 3.2607411432266233 \t mean correct 1.970972\t val loss 345.13307006835936 \t val KL 3.9012134075164795 \t val mean correct 1.978230\t took 19.411728 seconds\n",
      "training rollout len is 2. Beta: 64.0\n",
      "epoch 53 \t loss 675.089894 \t KL 3.6199264550209045 \t mean correct 1.970760\t val loss 357.34033935546876 \t val KL 4.243165016174316 \t val mean correct 1.976519\t took 19.742188 seconds\n",
      "training rollout len is 2. Beta: 64.0\n",
      "epoch 54 \t loss 666.401252 \t KL 3.652542926073074 \t mean correct 1.971694\t val loss 352.40544067382814 \t val KL 3.3988842964172363 \t val mean correct 1.977331\t took 19.653594 seconds\n",
      "training rollout len is 2. Beta: 64.0\n",
      "epoch 55 \t loss 646.392505 \t KL 3.2653419625759126 \t mean correct 1.971432\t val loss 362.15899780273435 \t val KL 3.110630750656128 \t val mean correct 1.976736\t took 19.728888 seconds\n",
      "training rollout len is 2. Beta: 64.0\n",
      "epoch 56 \t loss 665.617329 \t KL 3.5918715739250184 \t mean correct 1.971299\t val loss 348.4109356689453 \t val KL 3.2004404067993164 \t val mean correct 1.976955\t took 19.547282 seconds\n",
      "training rollout len is 2. Beta: 64.0\n",
      "epoch 57 \t loss 637.149059 \t KL 2.92157856464386 \t mean correct 1.970402\t val loss 354.05100830078123 \t val KL 3.8494203090667725 \t val mean correct 1.977036\t took 19.594030 seconds\n",
      "training rollout len is 2. Beta: 64.0\n",
      "epoch 58 \t loss 663.947176 \t KL 3.5256598246097566 \t mean correct 1.971186\t val loss 348.719853515625 \t val KL 3.9217710494995117 \t val mean correct 1.976991\t took 19.456310 seconds\n",
      "training rollout len is 2. Beta: 64.0\n",
      "epoch 59 \t loss 676.732373 \t KL 3.5463924396038053 \t mean correct 1.970673\t val loss 358.7275732421875 \t val KL 3.2570245265960693 \t val mean correct 1.977491\t took 19.630829 seconds\n",
      "training rollout len is 3. Beta: 64.0\n",
      "epoch 60 \t loss 732.876256 \t KL 3.125646166006724 \t mean correct 1.965657\t val loss 375.49498779296874 \t val KL 3.2992217540740967 \t val mean correct 1.976623\t took 27.064502 seconds\n",
      "training rollout len is 3. Beta: 64.0\n",
      "epoch 61 \t loss 742.359143 \t KL 3.3131279818216965 \t mean correct 1.965642\t val loss 390.98164916992187 \t val KL 3.32076358795166 \t val mean correct 1.976664\t took 26.985284 seconds\n",
      "training rollout len is 3. Beta: 64.0\n",
      "epoch 62 \t loss 794.338135 \t KL 4.1080841263135275 \t mean correct 1.965677\t val loss 368.77623168945314 \t val KL 3.6794493198394775 \t val mean correct 1.977348\t took 27.024196 seconds\n",
      "training rollout len is 3. Beta: 64.0\n",
      "epoch 63 \t loss 763.807902 \t KL 3.642370456854502 \t mean correct 1.965627\t val loss 391.5554357910156 \t val KL 4.33988618850708 \t val mean correct 1.977755\t took 27.075681 seconds\n",
      "training rollout len is 3. Beta: 64.0\n",
      "epoch 64 \t loss 745.748148 \t KL 3.240957667032876 \t mean correct 1.965287\t val loss 385.0033447265625 \t val KL 3.7064852714538574 \t val mean correct 1.977416\t took 27.215760 seconds\n",
      "training rollout len is 3. Beta: 64.0\n",
      "epoch 65 \t loss 779.100388 \t KL 4.053609001636505 \t mean correct 1.966368\t val loss 386.9458166503906 \t val KL 3.455634355545044 \t val mean correct 1.977270\t took 27.116244 seconds\n",
      "training rollout len is 3. Beta: 64.0\n",
      "epoch 66 \t loss 765.515499 \t KL 3.7658371893564837 \t mean correct 1.965940\t val loss 389.90014892578125 \t val KL 4.660678863525391 \t val mean correct 1.976644\t took 27.167408 seconds\n",
      "training rollout len is 3. Beta: 64.0\n",
      "epoch 67 \t loss 786.206926 \t KL 4.120804820060729 \t mean correct 1.966102\t val loss 355.31233642578127 \t val KL 3.793060302734375 \t val mean correct 1.977848\t took 27.108851 seconds\n",
      "training rollout len is 3. Beta: 64.0\n",
      "epoch 68 \t loss 795.795602 \t KL 4.031618944803872 \t mean correct 1.965171\t val loss 369.7323681640625 \t val KL 4.961526870727539 \t val mean correct 1.977772\t took 27.070545 seconds\n",
      "training rollout len is 3. Beta: 64.0\n",
      "epoch 69 \t loss 755.669243 \t KL 3.6113595533370972 \t mean correct 1.965962\t val loss 382.89378540039064 \t val KL 3.3255937099456787 \t val mean correct 1.977039\t took 26.975644 seconds\n",
      "training rollout len is 4. Beta: 64.0\n",
      "epoch 70 \t loss 856.775253 \t KL 3.7799536603689194 \t mean correct 1.960703\t val loss 377.22037109375 \t val KL 4.455796241760254 \t val mean correct 1.976892\t took 34.339065 seconds\n",
      "training rollout len is 4. Beta: 64.0\n",
      "epoch 71 \t loss 869.198571 \t KL 4.216055834889412 \t mean correct 1.961460\t val loss 370.140224609375 \t val KL 3.4692177772521973 \t val mean correct 1.977714\t took 34.493651 seconds\n",
      "training rollout len is 4. Beta: 64.0\n",
      "epoch 72 \t loss 866.310945 \t KL 4.008956789374351 \t mean correct 1.960951\t val loss 413.87826049804687 \t val KL 2.982692003250122 \t val mean correct 1.978144\t took 34.288174 seconds\n",
      "training rollout len is 4. Beta: 64.0\n",
      "epoch 73 \t loss 862.532264 \t KL 4.173536041378975 \t mean correct 1.961508\t val loss 399.38035888671874 \t val KL 4.106518268585205 \t val mean correct 1.977331\t took 34.268019 seconds\n",
      "training rollout len is 4. Beta: 64.0\n",
      "epoch 74 \t loss 880.382365 \t KL 4.366364610791206 \t mean correct 1.961692\t val loss 407.3464416503906 \t val KL 3.7345597743988037 \t val mean correct 1.976642\t took 34.542933 seconds\n",
      "training rollout len is 4. Beta: 64.0\n",
      "epoch 75 \t loss 859.667511 \t KL 3.9722108000516894 \t mean correct 1.960912\t val loss 397.3245483398438 \t val KL 4.267852306365967 \t val mean correct 1.977498\t took 34.217980 seconds\n",
      "val rollout loss: nan, val rollout mean correct: 1.8275992978852371, took 25.354722 seconds\n",
      "training rollout len is 4. Beta: 64.0\n",
      "epoch 76 \t loss 830.352094 \t KL 3.5295756036043167 \t mean correct 1.961024\t val loss 414.9554699707031 \t val KL 3.1748199462890625 \t val mean correct 1.977089\t took 34.442405 seconds\n",
      "training rollout len is 4. Beta: 64.0\n",
      "epoch 77 \t loss 859.814941 \t KL 3.9516360491514204 \t mean correct 1.961019\t val loss 409.6387451171875 \t val KL 3.144911527633667 \t val mean correct 1.976128\t took 34.376796 seconds\n",
      "training rollout len is 4. Beta: 64.0\n",
      "epoch 78 \t loss 831.050571 \t KL 3.5027308320999144 \t mean correct 1.961157\t val loss 400.87531005859375 \t val KL 3.999272584915161 \t val mean correct 1.977119\t took 34.292860 seconds\n",
      "training rollout len is 4. Beta: 64.0\n",
      "epoch 79 \t loss 822.165164 \t KL 3.4084073597192766 \t mean correct 1.961122\t val loss 395.4549230957031 \t val KL 4.4983391761779785 \t val mean correct 1.977578\t took 34.374617 seconds\n",
      "training rollout len is 5. Beta: 64.0\n",
      "epoch 80 \t loss 946.335254 \t KL 4.065378225326535 \t mean correct 1.956265\t val loss 456.94820068359377 \t val KL 3.8424501419067383 \t val mean correct 1.978037\t took 42.090585 seconds\n",
      "training rollout len is 5. Beta: 64.0\n",
      "epoch 81 \t loss 900.482276 \t KL 3.4192522969245918 \t mean correct 1.956468\t val loss 385.6043896484375 \t val KL 3.5712192058563232 \t val mean correct 1.978337\t took 42.127821 seconds\n",
      "training rollout len is 5. Beta: 64.0\n",
      "epoch 82 \t loss 826.207685 \t KL 2.4067514255046842 \t mean correct 1.956797\t val loss 418.33503173828126 \t val KL 2.719679832458496 \t val mean correct 1.976566\t took 42.114919 seconds\n",
      "training rollout len is 5. Beta: 64.0\n",
      "epoch 83 \t loss 885.380466 \t KL 3.054773700237275 \t mean correct 1.955859\t val loss 448.89469482421873 \t val KL 2.45668888092041 \t val mean correct 1.977494\t took 41.925301 seconds\n",
      "training rollout len is 5. Beta: 64.0\n",
      "epoch 84 \t loss 866.671143 \t KL 3.013933320999144 \t mean correct 1.956880\t val loss 433.8665563964844 \t val KL 2.193192720413208 \t val mean correct 1.977233\t took 42.352418 seconds\n",
      "training rollout len is 5. Beta: 64.0\n",
      "epoch 85 \t loss 890.210822 \t KL 3.28462702512741 \t mean correct 1.956448\t val loss 395.40751708984374 \t val KL 2.2770864963531494 \t val mean correct 1.977898\t took 42.455832 seconds\n",
      "training rollout len is 5. Beta: 64.0\n",
      "epoch 86 \t loss 871.460151 \t KL 2.942960332870482 \t mean correct 1.956367\t val loss 457.59697387695314 \t val KL 2.7209954261779785 \t val mean correct 1.977216\t took 42.009808 seconds\n",
      "training rollout len is 5. Beta: 64.0\n",
      "epoch 87 \t loss 867.584534 \t KL 3.034460537433623 \t mean correct 1.956781\t val loss 428.32703979492186 \t val KL 2.7057929039001465 \t val mean correct 1.976562\t took 41.982394 seconds\n",
      "training rollout len is 5. Beta: 64.0\n",
      "epoch 88 \t loss 878.822822 \t KL 3.0402803094387063 \t mean correct 1.956037\t val loss 421.67045654296874 \t val KL 2.4609057903289795 \t val mean correct 1.979078\t took 42.126212 seconds\n",
      "training rollout len is 5. Beta: 64.0\n",
      "epoch 89 \t loss 871.884519 \t KL 3.0202771205902104 \t mean correct 1.956224\t val loss 387.92978759765623 \t val KL 3.4848806858062744 \t val mean correct 1.977980\t took 41.979122 seconds\n",
      "training rollout len is 6. Beta: 64.0\n",
      "epoch 90 \t loss 973.946969 \t KL 3.481892385085423 \t mean correct 1.951808\t val loss 454.82366455078125 \t val KL 3.554062843322754 \t val mean correct 1.977045\t took 50.625857 seconds\n",
      "training rollout len is 6. Beta: 64.0\n",
      "epoch 91 \t loss 966.903308 \t KL 3.2124095904827126 \t mean correct 1.951034\t val loss 432.0938562011719 \t val KL 4.072185516357422 \t val mean correct 1.977848\t took 50.394463 seconds\n",
      "training rollout len is 6. Beta: 64.0\n",
      "epoch 92 \t loss 970.603387 \t KL 3.65970326979955 \t mean correct 1.952452\t val loss 454.935712890625 \t val KL 3.9677648544311523 \t val mean correct 1.976622\t took 50.410174 seconds\n",
      "training rollout len is 6. Beta: 64.0\n",
      "epoch 93 \t loss 960.259553 \t KL 3.3548545189698533 \t mean correct 1.951145\t val loss 377.3854614257813 \t val KL 4.0811357498168945 \t val mean correct 1.978137\t took 50.432697 seconds\n",
      "training rollout len is 6. Beta: 64.0\n",
      "epoch 94 \t loss 966.605076 \t KL 3.6512784071763362 \t mean correct 1.952063\t val loss 406.155029296875 \t val KL 3.9013826847076416 \t val mean correct 1.976698\t took 50.408990 seconds\n",
      "training rollout len is 6. Beta: 64.0\n",
      "epoch 95 \t loss 951.780614 \t KL 3.5368197143077835 \t mean correct 1.952381\t val loss 414.82060546875 \t val KL 3.769768238067627 \t val mean correct 1.977691\t took 50.431165 seconds\n",
      "training rollout len is 6. Beta: 64.0\n",
      "epoch 96 \t loss 952.235185 \t KL 3.687561167081198 \t mean correct 1.953035\t val loss 380.6390747070312 \t val KL 4.079091548919678 \t val mean correct 1.977411\t took 50.422377 seconds\n",
      "training rollout len is 6. Beta: 64.0\n",
      "epoch 97 \t loss 947.502269 \t KL 3.8093771811326342 \t mean correct 1.953356\t val loss 381.9846301269531 \t val KL 4.519263744354248 \t val mean correct 1.977881\t took 50.410088 seconds\n",
      "training rollout len is 6. Beta: 64.0\n",
      "epoch 98 \t loss 920.634583 \t KL 3.513547914028166 \t mean correct 1.953395\t val loss 405.3469348144531 \t val KL 5.103233337402344 \t val mean correct 1.975558\t took 50.412166 seconds\n",
      "training rollout len is 6. Beta: 64.0\n",
      "epoch 99 \t loss 948.571349 \t KL 3.596851753393809 \t mean correct 1.951476\t val loss 374.9910900878906 \t val KL 5.256844520568848 \t val mean correct 1.977162\t took 50.416822 seconds\n",
      "training rollout len is 7. Beta: 64.0\n",
      "epoch 100 \t loss 999.622688 \t KL 3.727867154053279 \t mean correct 1.948351\t val loss 378.55879150390626 \t val KL 4.877638816833496 \t val mean correct 1.976562\t took 57.351444 seconds\n",
      "val rollout loss: nan, val rollout mean correct: 1.8266257818813982, took 25.141286 seconds\n",
      "training rollout len is 7. Beta: 64.0\n",
      "epoch 101 \t loss 997.353054 \t KL 3.692088895525251 \t mean correct 1.947967\t val loss 361.85759765625 \t val KL 3.9683938026428223 \t val mean correct 1.977353\t took 57.572109 seconds\n",
      "training rollout len is 7. Beta: 64.0\n",
      "epoch 102 \t loss 988.810810 \t KL 3.712128872530802 \t mean correct 1.948801\t val loss 398.74291137695315 \t val KL 2.7893688678741455 \t val mean correct 1.977481\t took 57.602625 seconds\n",
      "training rollout len is 7. Beta: 64.0\n",
      "epoch 103 \t loss 974.622357 \t KL 3.364160044193267 \t mean correct 1.947719\t val loss 382.8373718261719 \t val KL 4.866745948791504 \t val mean correct 1.976973\t took 57.501811 seconds\n",
      "training rollout len is 7. Beta: 64.0\n",
      "epoch 104 \t loss 977.955902 \t KL 3.6772010389396126 \t mean correct 1.948803\t val loss 356.97304077148436 \t val KL 4.730661392211914 \t val mean correct 1.978120\t took 57.421427 seconds\n",
      "training rollout len is 7. Beta: 64.0\n",
      "epoch 105 \t loss 946.485514 \t KL 3.605246173483985 \t mean correct 1.950809\t val loss 361.01083984375 \t val KL 4.4108428955078125 \t val mean correct 1.977617\t took 57.246411 seconds\n",
      "training rollout len is 7. Beta: 64.0\n",
      "epoch 106 \t loss 922.543674 \t KL 3.628715665340424 \t mean correct 1.952568\t val loss 401.82007690429685 \t val KL 3.759627103805542 \t val mean correct 1.976587\t took 57.500539 seconds\n",
      "training rollout len is 7. Beta: 64.0\n",
      "epoch 107 \t loss 900.520884 \t KL 3.653706036635809 \t mean correct 1.954575\t val loss 354.0634069824219 \t val KL 2.7335243225097656 \t val mean correct 1.977745\t took 57.151929 seconds\n",
      "training rollout len is 7. Beta: 64.0\n",
      "epoch 108 \t loss 853.450754 \t KL 3.6136468802179587 \t mean correct 1.957859\t val loss 351.17034912109375 \t val KL 3.6426875591278076 \t val mean correct 1.978342\t took 57.181830 seconds\n",
      "training rollout len is 7. Beta: 64.0\n",
      "epoch 109 \t loss 831.544601 \t KL 3.609771901539395 \t mean correct 1.959540\t val loss 356.02469848632813 \t val KL 4.594994068145752 \t val mean correct 1.976711\t took 57.411326 seconds\n",
      "training rollout len is 8. Beta: 64.0\n",
      "epoch 110 \t loss 853.299741 \t KL 3.829408021867275 \t mean correct 1.959415\t val loss 369.1648095703125 \t val KL 4.45845890045166 \t val mean correct 1.976400\t took 64.782126 seconds\n",
      "training rollout len is 8. Beta: 64.0\n",
      "epoch 111 \t loss 823.225694 \t KL 3.7621508663892747 \t mean correct 1.961023\t val loss 343.96984619140625 \t val KL 4.490970611572266 \t val mean correct 1.977758\t took 64.710137 seconds\n",
      "training rollout len is 8. Beta: 64.0\n",
      "epoch 112 \t loss 810.099588 \t KL 3.8847270077466964 \t mean correct 1.962437\t val loss 334.02100952148436 \t val KL 3.6409499645233154 \t val mean correct 1.979031\t took 64.810290 seconds\n",
      "training rollout len is 8. Beta: 64.0\n",
      "epoch 113 \t loss 797.997503 \t KL 3.969636509269476 \t mean correct 1.963773\t val loss 333.264560546875 \t val KL 3.290449857711792 \t val mean correct 1.979078\t took 64.565890 seconds\n",
      "training rollout len is 8. Beta: 64.0\n",
      "epoch 114 \t loss 810.646539 \t KL 4.039277416169643 \t mean correct 1.963076\t val loss 346.8045971679687 \t val KL 3.973738431930542 \t val mean correct 1.977758\t took 64.674975 seconds\n",
      "training rollout len is 8. Beta: 64.0\n",
      "epoch 115 \t loss 812.927067 \t KL 4.18162682235241 \t mean correct 1.963395\t val loss 320.10406616210935 \t val KL 3.6903927326202393 \t val mean correct 1.979642\t took 64.717617 seconds\n",
      "training rollout len is 8. Beta: 64.0\n",
      "epoch 116 \t loss 779.044051 \t KL 3.943329378515482 \t mean correct 1.964709\t val loss 330.6025134277344 \t val KL 4.565425395965576 \t val mean correct 1.978587\t took 64.536704 seconds\n",
      "training rollout len is 8. Beta: 64.0\n",
      "epoch 117 \t loss 787.861106 \t KL 4.125396551191807 \t mean correct 1.965041\t val loss 328.70320678710937 \t val KL 3.8216476440429688 \t val mean correct 1.978748\t took 64.640372 seconds\n",
      "training rollout len is 8. Beta: 64.0\n",
      "epoch 118 \t loss 767.603175 \t KL 4.001139732301235 \t mean correct 1.965726\t val loss 326.44146118164065 \t val KL 5.656978607177734 \t val mean correct 1.979227\t took 64.872553 seconds\n",
      "training rollout len is 8. Beta: 64.0\n",
      "epoch 119 \t loss 775.374534 \t KL 4.181849740743637 \t mean correct 1.966218\t val loss 323.80912353515623 \t val KL 5.218495845794678 \t val mean correct 1.979344\t took 65.116579 seconds\n",
      "training rollout len is 9. Beta: 64.0\n",
      "epoch 120 \t loss 781.509970 \t KL 3.9331346382035144 \t mean correct 1.964678\t val loss 334.1175842285156 \t val KL 5.094382286071777 \t val mean correct 1.979236\t took 72.511927 seconds\n",
      "training rollout len is 9. Beta: 64.0\n",
      "epoch 121 \t loss 773.137468 \t KL 4.043098844952056 \t mean correct 1.965750\t val loss 332.44845092773437 \t val KL 4.294967174530029 \t val mean correct 1.979309\t took 72.316989 seconds\n",
      "training rollout len is 9. Beta: 64.0\n",
      "epoch 122 \t loss 760.227924 \t KL 4.122893368668025 \t mean correct 1.966929\t val loss 305.9982080078125 \t val KL 6.001380443572998 \t val mean correct 1.980458\t took 74.484671 seconds\n",
      "training rollout len is 9. Beta: 64.0\n",
      "epoch 123 \t loss 772.502828 \t KL 4.177943811151717 \t mean correct 1.966306\t val loss 308.2047497558594 \t val KL 4.408205032348633 \t val mean correct 1.980447\t took 72.411231 seconds\n",
      "training rollout len is 9. Beta: 64.0\n",
      "epoch 124 \t loss 760.647842 \t KL 4.135071295102438 \t mean correct 1.966958\t val loss 318.77754760742187 \t val KL 4.142354488372803 \t val mean correct 1.980023\t took 72.347496 seconds\n",
      "training rollout len is 9. Beta: 64.0\n",
      "epoch 125 \t loss 760.478999 \t KL 4.278229548401304 \t mean correct 1.967575\t val loss 329.09404907226565 \t val KL 4.865312576293945 \t val mean correct 1.979606\t took 72.639508 seconds\n",
      "val rollout loss: nan, val rollout mean correct: 1.9457150742103313, took 25.193013 seconds\n",
      "training rollout len is 9. Beta: 64.0\n",
      "epoch 126 \t loss 742.965513 \t KL 4.185020489957597 \t mean correct 1.968232\t val loss 312.21437194824216 \t val KL 4.36427116394043 \t val mean correct 1.980036\t took 72.425709 seconds\n",
      "training rollout len is 9. Beta: 64.0\n",
      "epoch 127 \t loss 760.988891 \t KL 4.267487896018558 \t mean correct 1.967518\t val loss 308.7816650390625 \t val KL 5.151121616363525 \t val mean correct 1.979812\t took 72.168351 seconds\n",
      "training rollout len is 9. Beta: 64.0\n",
      "epoch 128 \t loss 753.541314 \t KL 4.341004062493642 \t mean correct 1.968254\t val loss 304.4765832519531 \t val KL 4.361521244049072 \t val mean correct 1.980209\t took 72.309653 seconds\n",
      "training rollout len is 9. Beta: 64.0\n",
      "epoch 129 \t loss 731.250318 \t KL 4.183071385489571 \t mean correct 1.968884\t val loss 315.143779296875 \t val KL 3.7207775115966797 \t val mean correct 1.979727\t took 72.271847 seconds\n",
      "training rollout len is 10. Beta: 64.0\n",
      "epoch 130 \t loss 739.191924 \t KL 4.341031888723375 \t mean correct 1.969130\t val loss 311.09796630859375 \t val KL 4.498096942901611 \t val mean correct 1.979922\t took 79.910884 seconds\n",
      "training rollout len is 10. Beta: 64.0\n",
      "epoch 131 \t loss 749.153651 \t KL 4.179908758521081 \t mean correct 1.967868\t val loss 331.09923828125 \t val KL 4.530632972717285 \t val mean correct 1.979380\t took 79.936586 seconds\n",
      "training rollout len is 10. Beta: 64.0\n",
      "epoch 132 \t loss 749.161973 \t KL 4.258913872003556 \t mean correct 1.968210\t val loss 306.3139379882812 \t val KL 3.6615262031555176 \t val mean correct 1.980064\t took 79.955108 seconds\n",
      "training rollout len is 10. Beta: 64.0\n",
      "epoch 133 \t loss 733.635731 \t KL 4.200282264947893 \t mean correct 1.968968\t val loss 320.6580615234375 \t val KL 4.067329406738281 \t val mean correct 1.979470\t took 80.159314 seconds\n",
      "training rollout len is 10. Beta: 64.0\n",
      "epoch 134 \t loss 748.423204 \t KL 4.447467546463014 \t mean correct 1.969238\t val loss 311.47166015625 \t val KL 5.686050891876221 \t val mean correct 1.980161\t took 79.918347 seconds\n",
      "training rollout len is 10. Beta: 64.0\n",
      "epoch 135 \t loss 756.372819 \t KL 4.55080953645706 \t mean correct 1.969030\t val loss 304.20626037597657 \t val KL 4.76294469833374 \t val mean correct 1.980483\t took 79.737298 seconds\n",
      "training rollout len is 10. Beta: 64.0\n",
      "epoch 136 \t loss 745.254064 \t KL 4.518499865293503 \t mean correct 1.969669\t val loss 303.1224890136719 \t val KL 4.69577693939209 \t val mean correct 1.980602\t took 79.901750 seconds\n",
      "training rollout len is 10. Beta: 64.0\n",
      "epoch 137 \t loss 747.177740 \t KL 4.630153610467909 \t mean correct 1.969920\t val loss 317.51989990234375 \t val KL 4.0772385597229 \t val mean correct 1.979566\t took 79.697144 seconds\n",
      "training rollout len is 10. Beta: 64.0\n",
      "epoch 138 \t loss 739.660577 \t KL 4.403360576868059 \t mean correct 1.969586\t val loss 306.0018322753906 \t val KL 4.463028907775879 \t val mean correct 1.980272\t took 79.950703 seconds\n",
      "training rollout len is 10. Beta: 64.0\n",
      "epoch 139 \t loss 737.866997 \t KL 4.539778030633927 \t mean correct 1.970117\t val loss 293.9629187011719 \t val KL 4.103736400604248 \t val mean correct 1.980787\t took 80.140521 seconds\n",
      "training rollout len is 11. Beta: 64.0\n",
      "epoch 140 \t loss 730.542043 \t KL 4.353012679706921 \t mean correct 1.970091\t val loss 316.8554370117188 \t val KL 2.8534693717956543 \t val mean correct 1.980223\t took 87.284460 seconds\n",
      "training rollout len is 11. Beta: 64.0\n",
      "epoch 141 \t loss 735.175804 \t KL 4.431757323958658 \t mean correct 1.969996\t val loss 309.7464202880859 \t val KL 3.8957324028015137 \t val mean correct 1.980145\t took 87.299270 seconds\n",
      "training rollout len is 11. Beta: 64.0\n",
      "epoch 142 \t loss 721.538164 \t KL 4.214075881351124 \t mean correct 1.969982\t val loss 304.8082177734375 \t val KL 4.025666236877441 \t val mean correct 1.980314\t took 87.215051 seconds\n",
      "training rollout len is 11. Beta: 64.0\n",
      "epoch 143 \t loss 741.549636 \t KL 4.55736479022286 \t mean correct 1.970101\t val loss 307.3102453613281 \t val KL 3.233905553817749 \t val mean correct 1.980334\t took 87.231129 seconds\n",
      "training rollout len is 11. Beta: 64.0\n",
      "epoch 144 \t loss 729.839594 \t KL 4.3686160518906325 \t mean correct 1.970042\t val loss 309.16670288085936 \t val KL 5.527931213378906 \t val mean correct 1.979794\t took 87.347102 seconds\n",
      "training rollout len is 11. Beta: 64.0\n",
      "epoch 145 \t loss 736.899533 \t KL 4.685685890804635 \t mean correct 1.970846\t val loss 297.2632666015625 \t val KL 3.541914701461792 \t val mean correct 1.980859\t took 87.443553 seconds\n",
      "training rollout len is 11. Beta: 64.0\n",
      "epoch 146 \t loss 724.324392 \t KL 4.461545126438141 \t mean correct 1.970698\t val loss 312.4635888671875 \t val KL 3.851577043533325 \t val mean correct 1.979556\t took 87.471651 seconds\n",
      "training rollout len is 11. Beta: 64.0\n",
      "epoch 147 \t loss 712.311859 \t KL 4.39008030089465 \t mean correct 1.971182\t val loss 302.78509765625 \t val KL 3.5670735836029053 \t val mean correct 1.980072\t took 87.207087 seconds\n",
      "training rollout len is 11. Beta: 64.0\n",
      "epoch 148 \t loss 699.390271 \t KL 4.355955584157597 \t mean correct 1.971883\t val loss 303.8641833496094 \t val KL 4.859329700469971 \t val mean correct 1.979942\t took 87.364216 seconds\n",
      "training rollout len is 11. Beta: 64.0\n",
      "epoch 149 \t loss 715.988786 \t KL 4.504135895642367 \t mean correct 1.971293\t val loss 297.2808728027344 \t val KL 4.3521952629089355 \t val mean correct 1.980505\t took 87.224097 seconds\n",
      "training rollout len is 12. Beta: 64.0\n",
      "epoch 150 \t loss 713.911078 \t KL 4.414368924498557 \t mean correct 1.971276\t val loss 303.5105810546875 \t val KL 3.7285704612731934 \t val mean correct 1.980647\t took 95.085901 seconds\n",
      "val rollout loss: nan, val rollout mean correct: 1.965642968539534, took 25.349342 seconds\n",
      "training rollout len is 12. Beta: 64.0\n",
      "epoch 151 \t loss 710.418819 \t KL 4.398797722260157 \t mean correct 1.971496\t val loss 286.30759887695314 \t val KL 4.54810094833374 \t val mean correct 1.981414\t took 94.849527 seconds\n",
      "training rollout len is 12. Beta: 64.0\n",
      "epoch 152 \t loss 716.652253 \t KL 4.663151670694352 \t mean correct 1.972048\t val loss 291.6601135253906 \t val KL 4.743757724761963 \t val mean correct 1.981102\t took 95.607460 seconds\n",
      "training rollout len is 12. Beta: 64.0\n",
      "epoch 153 \t loss 728.430592 \t KL 4.619985092083613 \t mean correct 1.971046\t val loss 287.8558544921875 \t val KL 4.036938190460205 \t val mean correct 1.981778\t took 95.405890 seconds\n",
      "training rollout len is 12. Beta: 64.0\n",
      "epoch 154 \t loss 700.564220 \t KL 4.525672570069631 \t mean correct 1.972417\t val loss 290.7390380859375 \t val KL 4.162302017211914 \t val mean correct 1.980847\t took 95.737671 seconds\n",
      "training rollout len is 12. Beta: 64.0\n",
      "epoch 155 \t loss 715.852643 \t KL 4.590208195447921 \t mean correct 1.971732\t val loss 301.8637634277344 \t val KL 5.2791666984558105 \t val mean correct 1.980817\t took 95.126680 seconds\n",
      "training rollout len is 12. Beta: 64.0\n",
      "epoch 156 \t loss 702.722499 \t KL 4.586483439803125 \t mean correct 1.972498\t val loss 293.1953393554688 \t val KL 4.756773471832275 \t val mean correct 1.980948\t took 95.025127 seconds\n",
      "training rollout len is 12. Beta: 64.0\n",
      "epoch 157 \t loss 702.384439 \t KL 4.3820217470328 \t mean correct 1.971698\t val loss 295.73799682617187 \t val KL 4.330549240112305 \t val mean correct 1.980798\t took 95.047293 seconds\n",
      "training rollout len is 12. Beta: 64.0\n",
      "epoch 158 \t loss 703.873917 \t KL 4.636534203489623 \t mean correct 1.972773\t val loss 287.1253125 \t val KL 4.506112575531006 \t val mean correct 1.981227\t took 98.802243 seconds\n",
      "training rollout len is 12. Beta: 64.0\n",
      "epoch 159 \t loss 704.246796 \t KL 4.506360159715016 \t mean correct 1.972255\t val loss 285.36814758300784 \t val KL 4.467097759246826 \t val mean correct 1.981677\t took 95.180235 seconds\n",
      "training rollout len is 13. Beta: 64.0\n",
      "epoch 160 \t loss 714.979694 \t KL 4.62582022648591 \t mean correct 1.971924\t val loss 293.41975830078127 \t val KL 5.177165985107422 \t val mean correct 1.980687\t took 102.769464 seconds\n",
      "training rollout len is 13. Beta: 64.0\n",
      "epoch 161 \t loss 709.309355 \t KL 4.497082495322594 \t mean correct 1.971865\t val loss 290.59034545898436 \t val KL 4.732914447784424 \t val mean correct 1.981277\t took 102.228467 seconds\n",
      "training rollout len is 13. Beta: 64.0\n",
      "epoch 162 \t loss 715.513722 \t KL 4.68979626912337 \t mean correct 1.972063\t val loss 293.9786706542969 \t val KL 5.0576558113098145 \t val mean correct 1.981177\t took 102.271751 seconds\n",
      "training rollout len is 13. Beta: 64.0\n",
      "epoch 163 \t loss 698.603649 \t KL 4.2791842166277085 \t mean correct 1.972022\t val loss 284.16048095703127 \t val KL 4.602059364318848 \t val mean correct 1.981447\t took 102.400429 seconds\n",
      "training rollout len is 13. Beta: 64.0\n",
      "epoch 164 \t loss 694.283610 \t KL 4.503149186647855 \t mean correct 1.972730\t val loss 293.6884655761719 \t val KL 4.449407577514648 \t val mean correct 1.980708\t took 102.310279 seconds\n",
      "training rollout len is 13. Beta: 64.0\n",
      "epoch 165 \t loss 702.546211 \t KL 4.6227551984787 \t mean correct 1.972658\t val loss 295.54610107421877 \t val KL 3.9898881912231445 \t val mean correct 1.980998\t took 102.326709 seconds\n",
      "training rollout len is 13. Beta: 64.0\n",
      "epoch 166 \t loss 692.445544 \t KL 4.485318370599013 \t mean correct 1.972942\t val loss 295.89632080078127 \t val KL 4.66391134262085 \t val mean correct 1.980841\t took 103.541471 seconds\n",
      "training rollout len is 13. Beta: 64.0\n",
      "epoch 167 \t loss 698.712320 \t KL 4.543870006524602 \t mean correct 1.972614\t val loss 290.48619873046874 \t val KL 4.797389507293701 \t val mean correct 1.981053\t took 102.140723 seconds\n",
      "training rollout len is 13. Beta: 64.0\n",
      "epoch 168 \t loss 701.208727 \t KL 4.638937356288617 \t mean correct 1.972863\t val loss 299.3682666015625 \t val KL 5.295055389404297 \t val mean correct 1.980694\t took 102.600074 seconds\n",
      "training rollout len is 13. Beta: 64.0\n",
      "epoch 169 \t loss 696.575385 \t KL 4.567775236276479 \t mean correct 1.972785\t val loss 284.57326171875 \t val KL 3.325124502182007 \t val mean correct 1.981642\t took 102.320709 seconds\n",
      "training rollout len is 14. Beta: 64.0\n",
      "epoch 170 \t loss 697.174096 \t KL 4.664576078823634 \t mean correct 1.973211\t val loss 296.7808508300781 \t val KL 5.314983367919922 \t val mean correct 1.980745\t took 109.939670 seconds\n",
      "training rollout len is 14. Beta: 64.0\n",
      "epoch 171 \t loss 686.703325 \t KL 4.441618948493684 \t mean correct 1.972964\t val loss 296.5632800292969 \t val KL 5.426346302032471 \t val mean correct 1.980912\t took 109.864247 seconds\n",
      "training rollout len is 14. Beta: 64.0\n",
      "epoch 172 \t loss 712.430381 \t KL 4.775032461370739 \t mean correct 1.972670\t val loss 306.6163525390625 \t val KL 4.9006123542785645 \t val mean correct 1.979583\t took 110.015839 seconds\n",
      "training rollout len is 14. Beta: 64.0\n",
      "epoch 173 \t loss 688.435282 \t KL 4.585691074473516 \t mean correct 1.973377\t val loss 301.6373986816406 \t val KL 4.394215106964111 \t val mean correct 1.980320\t took 109.984466 seconds\n",
      "training rollout len is 14. Beta: 64.0\n",
      "epoch 174 \t loss 687.970915 \t KL 4.516693606206348 \t mean correct 1.973169\t val loss 300.2636340332031 \t val KL 5.213338375091553 \t val mean correct 1.980567\t took 110.734231 seconds\n",
      "training rollout len is 14. Beta: 64.0\n",
      "epoch 175 \t loss 690.371520 \t KL 4.610883179903031 \t mean correct 1.973317\t val loss 278.8585888671875 \t val KL 3.552305221557617 \t val mean correct 1.981764\t took 109.847880 seconds\n",
      "val rollout loss: nan, val rollout mean correct: 1.9683210122996362, took 25.265596 seconds\n",
      "training rollout len is 14. Beta: 64.0\n",
      "epoch 176 \t loss 688.142940 \t KL 4.537219712904521 \t mean correct 1.973195\t val loss 293.1084045410156 \t val KL 4.7957587242126465 \t val mean correct 1.981075\t took 110.597019 seconds\n",
      "training rollout len is 14. Beta: 64.0\n",
      "epoch 177 \t loss 685.707167 \t KL 4.566836797850474 \t mean correct 1.973478\t val loss 297.9814086914063 \t val KL 4.5673065185546875 \t val mean correct 1.980306\t took 110.323770 seconds\n",
      "training rollout len is 14. Beta: 64.0\n"
     ]
    }
   ],
   "source": [
    "!python train_model.py cell_dynamics_EPNS \\\n",
    "    --data_directory=../simulations/morpheus-twocell \\\n",
    "    --save_path=../models/morpheus-twocell-800.pt \\\n",
    "    --device=cuda"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "PyTorch v2.x (Machine Learning)",
   "language": "python",
   "name": "pytorch_v2"
  },
  "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}