From ac2c4bb4c5f80de44d2d1d26efdaa77c756dafa4 Mon Sep 17 00:00:00 2001 From: bxdd Date: Mon, 20 Nov 2023 16:12:23 +0800 Subject: [PATCH] [FIX] fix import bugs, and update setup --- learnware/specification/system/__init__.py | 2 +- setup.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/learnware/specification/system/__init__.py b/learnware/specification/system/__init__.py index dea4a19..82fbe3f 100644 --- a/learnware/specification/system/__init__.py +++ b/learnware/specification/system/__init__.py @@ -8,4 +8,4 @@ if not is_torch_available(verbose=False): HeteroMapTableSpecification = None logger.warning("HeteroMapTableSpecification is skipped because torch is not installed!") else: - from .rkme import HeteroMapTableSpecification + from .hetero_table import HeteroMapTableSpecification diff --git a/setup.py b/setup.py index a3f25b5..2145a04 100644 --- a/setup.py +++ b/setup.py @@ -107,8 +107,9 @@ if __name__ == "__main__": "pre-commit", ], "full": [ - "torch>=2.0.0", - "torchvision>=0.15.1", + # The default full requirements for learnware package + "torch==2.1.0", + "torchvision==0.16.0", "torch-optimizer>=0.3.0", "lightgbm>=3.3.0", "sentence_transformers>=2.2.2",