The tensorflow errors are annoying. I hate them. We should convert the nbs to pytorch and hide the complexity where applicable.