This is a slightly modified version of torch_cg from https://github.com/sbarratt/torch_cg
