ATen/cuda/CUDAContext.h: No such file or directory

I’m trying to compile this project. (https://github.com/NVIDIA/flownet2-pytorch). They used Pytorch 0.4.1. I’m also using Pytorch 0.4.1. And I get this error: fatal error: ATen/cuda/CUDAContext.h: No such file or directory. I tried with Pytorch 1.0.0, it can compile successfully. So it seems this file is introduced in newer Pytorch version. How can I have this file in Pytorch 0.4.1? Thank you!