Samuel Thomas commited on
Commit
9c250a6
·
1 Parent(s): 0d26ceb
Files changed (1) hide show
  1. tools.py +3 -3
tools.py CHANGED
@@ -2257,7 +2257,7 @@ def create_llm_pipeline():
2257
  return pipeline(
2258
  "text-generation",
2259
  model=model_id,
2260
- device_map="auto",
2261
  torch_dtype=torch.float16,
2262
  max_new_tokens=1024,
2263
  temperature=0.3,
@@ -3062,8 +3062,8 @@ Current date: {current_date_str}
3062
 
3063
  # Cleanup
3064
  if result.get("done"):
3065
- torch.cuda.empty_cache()
3066
- torch.cuda.ipc_collect()
3067
  gc.collect()
3068
  print("🧹 Released GPU memory after completion")
3069
 
 
2257
  return pipeline(
2258
  "text-generation",
2259
  model=model_id,
2260
+ device_map="cpu",
2261
  torch_dtype=torch.float16,
2262
  max_new_tokens=1024,
2263
  temperature=0.3,
 
3062
 
3063
  # Cleanup
3064
  if result.get("done"):
3065
+ #torch.cuda.empty_cache()
3066
+ #torch.cuda.ipc_collect()
3067
  gc.collect()
3068
  print("🧹 Released GPU memory after completion")
3069