diff --git a/torch_xla/torch_xla.py b/torch_xla/torch_xla.py
index 5f79fc27f6c..ef8e30667da 100644
--- a/torch_xla/torch_xla.py
+++ b/torch_xla/torch_xla.py
@@ -1,4 +1,6 @@
+import contextlib
 from typing import List
+
 import torch
 import torch_xla
 import torch_xla.core.xla_model as xm