diff --git a/torch/csrc/jit/script/final_returns.cpp b/torch/csrc/jit/script/final_returns.cpp
index 2671ac273f6e..47534fefe823 100644
--- a/torch/csrc/jit/script/final_returns.cpp
+++ b/torch/csrc/jit/script/final_returns.cpp
@@ -41,7 +41,7 @@ ReturnInfo makeReturnsFinal(
     const SourceRange& range,
     at::ArrayRef<TreeRef> stmts,
     bool return_none) {
-  std::vector<TreeRef> changed;
+  at::SmallVector<TreeRef, 4> changed;
   changed.reserve(stmts.size());
   for (size_t i = 0; i < stmts.size(); ++i) {
     const TreeRef& stmt = stmts[i];
diff --git a/torch/csrc/jit/script/tree.h b/torch/csrc/jit/script/tree.h
index d0f65f7251c2..0a48b6434f47 100644
--- a/torch/csrc/jit/script/tree.h
+++ b/torch/csrc/jit/script/tree.h
@@ -6,6 +6,7 @@
 #include <vector>
 
 #include <torch/csrc/jit/script/lexer.h>
+#include <c10/util/SmallVector.h>
 
 namespace torch {
 namespace jit {
@@ -29,7 +30,7 @@ namespace script {
 
 struct Tree;
 using TreeRef = std::shared_ptr<Tree>;
-using TreeList = std::vector<TreeRef>;
+using TreeList = at::SmallVector<TreeRef, 4>;
 
 static const TreeList empty_trees = {};