Conrad Lippert-Zajaczkowski
commited on
Commit
·
c3b6111
1
Parent(s):
6c0a1e0
run on A100
Browse files- handler.py +1 -2
handler.py
CHANGED
@@ -25,7 +25,6 @@ class EndpointHandler:
|
|
25 |
"/repository",
|
26 |
device_map="auto",
|
27 |
torch_dtype=dtype,
|
28 |
-
offload_folder="offload",
|
29 |
local_files_only=True
|
30 |
)
|
31 |
gpu_info1 = nvmlDeviceGetMemoryInfo(gpu_h1)
|
@@ -52,7 +51,7 @@ class EndpointHandler:
|
|
52 |
top_k=10,
|
53 |
num_return_sequences=1,
|
54 |
eos_token_id=self.tokenizer.eos_token_id,
|
55 |
-
max_length=256
|
56 |
)
|
57 |
# postprocess the prediction
|
58 |
return prediction
|
|
|
25 |
"/repository",
|
26 |
device_map="auto",
|
27 |
torch_dtype=dtype,
|
|
|
28 |
local_files_only=True
|
29 |
)
|
30 |
gpu_info1 = nvmlDeviceGetMemoryInfo(gpu_h1)
|
|
|
51 |
top_k=10,
|
52 |
num_return_sequences=1,
|
53 |
eos_token_id=self.tokenizer.eos_token_id,
|
54 |
+
max_length=256
|
55 |
)
|
56 |
# postprocess the prediction
|
57 |
return prediction
|