diff --git a/deepmd/pt/model/task/property.py b/deepmd/pt/model/task/property.py
index c5d0285caf..233cd3ddb2 100644
--- a/deepmd/pt/model/task/property.py
+++ b/deepmd/pt/model/task/property.py
@@ -2,6 +2,7 @@
 import logging
 from typing import (
     Optional,
+    Union
 )
 
 import torch