In LayerNMF in vis/reward_vis.py. Also deal with places that call that. Because we don't use its tensorness!