Add disk cache slider

2025-06-05 21:59:24 +02:00 · 2022-06-19 22:53:30 -04:00
parent efed44ac8d
commit 47a58a36b8
4 changed files with 63 additions and 18 deletions
--- a/breakmodel.py
+++ b/breakmodel.py
@ -230,7 +230,8 @@ logger = logging.get_logger(__name__)

 breakmodel = True
 gpu_blocks = []
-primary_device = 0
+disk_blocks = 0
+primary_device = 0 if torch.cuda.device_count() > 0 else "cpu"


 # Copied from transformers.models.bart.modeling_bart._expand_mask