Merge branch 'develop' into tensor_parallel

dice-group · Nov 29, 2024 · e0ee128 · e0ee128
2 parents f1b263b + 58aa98c
commit e0ee128
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -1,3 +1,5 @@
+[![Downloads](https://static.pepy.tech/badge/dicee)](https://pepy.tech/project/dicee)
+[![Downloads](https://img.shields.io/pypi/dm/dicee)](https://pypi.org/project/dicee/)
 [![Coverage](https://img.shields.io/badge/coverage-54%25-green)](https://dice-group.github.io/dice-embeddings/usage/main.html#coverage-report)
 [![Pypi](https://img.shields.io/badge/pypi-0.1.4-blue)](https://pypi.org/project/dicee/0.1.4/)
 [![Docs](https://img.shields.io/badge/documentation-0.1.4-yellow)](https://dice-group.github.io/dice-embeddings/index.html)

diff --git a/dicee/trainer/model_parallelism.py b/dicee/trainer/model_parallelism.py
@@ -87,7 +87,7 @@ def increase_batch_size_until_cuda_out_of_memory(ensemble_model, train_loader, b
                     return batch_sizes_and_mem_usages,True
 
         except torch.OutOfMemoryError as e:
-            print(f"torch.OutOfMemoryError caught! {e}")
+            print(f"torch.OutOfMemoryError caught! {e}\n\n")
             return batch_sizes_and_mem_usages, False
 
     history_batch_sizes_and_mem_usages=[]

diff --git a/examples/multi_hop_query_answering/benchmarking.py b/examples/multi_hop_query_answering/benchmarking.py
@@ -31,7 +31,7 @@
     args = Namespace()
     args.model = kge_name
     args.scoring_technique = "KvsAll"
-    args.path_dataset_folder = "KGs/UMLS"
+    args.dataset_dir = "KGs/UMLS"
     args.num_epochs = 20
     args.batch_size = 1024
     args.lr = 0.1