diff --git a/learnware/utils/gpu.py b/learnware/utils/gpu.py index 95fbfe1..395d979 100644 --- a/learnware/utils/gpu.py +++ b/learnware/utils/gpu.py @@ -44,3 +44,17 @@ def choose_device(cuda_idx=-1): else: device = torch.device("cuda:0") return device + + +class CudaManager: + def __init__(self): + pass + + def get_specification_cuda_idx(self): + pass + + def get_cpu_idx(self): + pass + + def get_random_cuda_idx(self): + pass diff --git a/setup.py b/setup.py index 705000d..b6aea9a 100644 --- a/setup.py +++ b/setup.py @@ -69,6 +69,7 @@ REQUIRED = [ "transformers>=4.34.1", "portalocker>=2.0.0", "qpsolvers[clarabel]>=4.0.1", + "geatpy>=2.7.0;python_version<='3.11'", ] here = os.path.abspath(os.path.dirname(__file__)) @@ -105,6 +106,13 @@ if __name__ == "__main__": "black==23.1.0", "pre-commit", ], + "full": [ + "torch>=1.11.0", + "torchvision>=0.15.1", + "torch-optimizer>=0.3.0", + "lightgbm>=3.3.0", + "sentence_transformers>=2.2.2", + ], }, classifiers=[ "Intended Audience :: Science/Research",