Fix for new version

Lightning-Universe · Jun 15, 2022 · 4320245 · 4320245
1 parent 051f67b
commit 4320245
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tests/test_hp_space_generator.py b/tests/test_hp_space_generator.py
@@ -1,8 +1,9 @@
 from hp_space_generator import HPSpaceGenerator, RandomSearchStrategy, GridSearchStrategy
 
-from lightning import LightningWork
+import lightning as L
 
-class DoSomethingExtra(LightningWork):
+
+class DoSomethingExtra(L.LightningWork):
     def __init__(self):
         super().__init__(run_once=True)
         self.hp_list = []