Conrad Lippert-Zajaczkowski commited on
Commit
c3b6111
1 Parent(s): 6c0a1e0

run on A100

Browse files
Files changed (1) hide show
  1. 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