diff --git a/climax/main/lr_range_test.py b/climax/main/lr_range_test.py
index 99e716e6da74cbf1f1fd3c5b4d5db4b327c4c8ff..1bdaa16045b05a7fefa76968a888b7490027a6a5 100644
--- a/climax/main/lr_range_test.py
+++ b/climax/main/lr_range_test.py
@@ -11,6 +11,7 @@ from datetime import timedelta
 from logging.config import dictConfig
 
 # externals
+import torch
 import xarray as xr
 from sklearn.model_selection import train_test_split
 from torch.utils.data import DataLoader