diff --git a/test/test_distribution_utils/test_stabilize_derivative.py b/test/test_distribution_utils/test_stabilize_derivative.py
index f2c71562..6377a67e 100644
--- a/test/test_distribution_utils/test_stabilize_derivative.py
+++ b/test/test_distribution_utils/test_stabilize_derivative.py
@@ -1,5 +1,6 @@
 from ..utils import BaseTestClass
-from xgboostlss.model import *
+import torch
+
 
 
 class TestClass(BaseTestClass):