Imported from https://github.com/Navhkrin/Bts-PyTorch