Browse Source

Merge branch 'fix-cache' of github.com:bigscience-workshop/petals into fix-cache

justheuristic 3 years ago
parent
commit
a9af1b4cf5
1 changed files with 1 additions and 1 deletions
  1. 1 1
      src/server/handler.py

+ 1 - 1
src/server/handler.py

@@ -51,7 +51,7 @@ class TransformerConnectionHandler(ConnectionHandler):
             if not requested_uids:
             if not requested_uids:
                 raise ValueError("User must specify at least one block for inference, but got none")
                 raise ValueError("User must specify at least one block for inference, but got none")
             assert isinstance(max_length, int), f"rpc_inference metadata must contain int seq_length, got {max_length}"
             assert isinstance(max_length, int), f"rpc_inference metadata must contain int seq_length, got {max_length}"
-            if max_length not in range(0, self.inference_max_length):
+            if not 0 <= max_length <= self.inference_max_length:
                 raise ValueError(f"Cannot allocate KV cache for {max_length} tokens, max = {self.inference_max_length}")
                 raise ValueError(f"Cannot allocate KV cache for {max_length} tokens, max = {self.inference_max_length}")
 
 
             batch_size = request.tensors[0].size[0] if request.tensors else 1
             batch_size = request.tensors[0].size[0] if request.tensors else 1