diff --git a/opencompass/partitioners/size.py b/opencompass/partitioners/size.py
index 56659ca0..3efba82a 100644
--- a/opencompass/partitioners/size.py
+++ b/opencompass/partitioners/size.py
@@ -1,6 +1,7 @@
 import copy
 import math
 import os.path as osp
+import torch
 from fnmatch import fnmatch
 from typing import Dict, List, Optional, Tuple, Union