Skip to content

Commit

Permalink
Fix compatibility
Browse files Browse the repository at this point in the history
  • Loading branch information
saharNooby committed Jun 21, 2023
1 parent 3767697 commit b92fbaa
Showing 1 changed file with 4 additions and 0 deletions.
4 changes: 4 additions & 0 deletions rwkv/rwkv_cpp_model.py
Expand Up @@ -15,6 +15,7 @@ def __init__(
model_path: str,
thread_count: int = max(1, multiprocessing.cpu_count() // 2),
gpu_layer_count: int = 0,
**kwargs
):
"""
Loads the model and prepares it for inference.
Expand All @@ -32,6 +33,9 @@ def __init__(
Count of layers to offload onto the GPU, must be >= 0.
"""

if 'gpu_layers_count' in kwargs:
gpu_layer_count = kwargs['gpu_layers_count']

assert os.path.isfile(model_path), f'{model_path} is not a file'
assert thread_count > 0, 'Thread count must be > 0'
assert gpu_layer_count >= 0, 'GPU layer count must be >= 0'
Expand Down

0 comments on commit b92fbaa

Please sign in to comment.