update readme

code clean up

change loss dec

change likelihood to rank 0

change likelihood

refactor filename

add more prints

add print steps

training only no plotting

remove make gif function

remove load from checkpoint

change plotting to headless Agg

remove data limit of 15000

add save model after training

update code