diff --git a/shape_inference.h b/shape_inference.h
index da2a2ed..ede73d9 100644
--- a/shape_inference.h
+++ b/shape_inference.h
@@ -1,6 +1,6 @@
 #pragma once
 
-#include <aten/src/ATen/Context.h>
+#include <ATen/Context.h>
 #include <torch/csrc/jit/ir/ir.h>
 
 namespace torch {