diff --git a/torch_xla/csrc/runtime/tensor_source.h b/torch_xla/csrc/runtime/tensor_source.h
index 38697942e5b..ffba51d80be 100644
--- a/torch_xla/csrc/runtime/tensor_source.h
+++ b/torch_xla/csrc/runtime/tensor_source.h
@@ -1,9 +1,14 @@
 #ifndef XLA_CLIENT_TENSOR_SOURCE_H_
 #define XLA_CLIENT_TENSOR_SOURCE_H_
 
+#include <vector>
+
+#include <ATen/Tensor.h>
+
 #include "torch_xla/csrc/runtime/debug_macros.h"
 #include "xla/literal.h"
 #include "xla/shape.h"
+#include "xla/shape_util.h"
 
 namespace torch_xla {
 namespace runtime {