Fix for new version

Lightning-Universe · Jun 15, 2022 · 051f67b · 051f67b
1 parent 1ba45af
commit 051f67b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/hp_space_generator/component.py b/hp_space_generator/component.py
@@ -1,7 +1,7 @@
-from lightning import LightningFlow
+import lightning as L
 
 
-class HPSpaceGenerator(LightningFlow):
+class HPSpaceGenerator(L.LightningFlow):
     """
     The HP Engine Component is used to suggest a list of configurations (hyper-parameters) to run with some config
     from the user for any task.

diff --git a/hp_space_generator/search_strategy.py b/hp_space_generator/search_strategy.py
@@ -1,10 +1,10 @@
 from typing import Dict, Any
 
-from lightning import LightningWork
+import lightning as L
 
 
 # TODO: Maybe create a method which makes sure that the given format for the hpe Config Dictionary is correct
-class SearchStrategy(LightningWork):
+class SearchStrategy(L.LightningWork):
     def __init__(self, should_preprocess=True, *args, **kwargs):
         super().__init__(run_once=True, parallel=True, *args, **kwargs)
         self.should_preprocess = should_preprocess