There are a lot of things in the PyTorch notebook that make this assignment a lot easier. They should be backported to the TF version. One thing to note is that `n_lives` exists in the PyTorch version, but not in TF.