diff --git a/learnware/config.py b/learnware/config.py index d7b056f..13e212f 100644 --- a/learnware/config.py +++ b/learnware/config.py @@ -119,6 +119,7 @@ semantic_config = { _DEFAULT_CONFIG = { "root_path": ROOT_DIRPATH, "package_path": PACKAGE_DIRPATH, + "database_path": DATABASE_PATH, "stdout_path": STDOUT_PATH, "cache_path": CACHE_PATH, "logging_level": logging.INFO, diff --git a/learnware/market/heterogeneous/organizer/hetero_map/__init__.py b/learnware/market/heterogeneous/organizer/hetero_map/__init__.py index fdb2137..2e31195 100644 --- a/learnware/market/heterogeneous/organizer/hetero_map/__init__.py +++ b/learnware/market/heterogeneous/organizer/hetero_map/__init__.py @@ -80,8 +80,8 @@ class HeteroMap(nn.Module): """ super(HeteroMap, self).__init__() - self.cuda_idx = allocate_cuda_idx() if cuda_idx is None else cuda_idx - self.device = choose_device(self.cuda_idx) + cuda_idx = allocate_cuda_idx() if cuda_idx is None else cuda_idx + device = choose_device(cuda_idx) self.model_args = { "num_partition": num_partition, "overlap_ratio": overlap_ratio, @@ -105,7 +105,7 @@ class HeteroMap(nn.Module): pad_token_id=feature_tokenizer.pad_token_id, hidden_dim=hidden_dim, hidden_dropout_prob=hidden_dropout_prob, - device=self.device, + device=device, ) self.encoder = TransformerMultiLayer( @@ -127,7 +127,7 @@ class HeteroMap(nn.Module): self.base_temperature = base_temperature self.num_partition = num_partition self.overlap_ratio = overlap_ratio - self.to(self.device) + self.to(device) def to(self, device: Union[str, torch.device]): """Moves the model and all its submodules to the specified device diff --git a/learnware/utils/gpu.py b/learnware/utils/gpu.py index cb856d4..bc19979 100644 --- a/learnware/utils/gpu.py +++ b/learnware/utils/gpu.py @@ -47,7 +47,7 @@ def choose_device(cuda_idx=-1): return device -def allocate_cuda_idx(self): +def allocate_cuda_idx(): if is_torch_available(verbose=False): import torch