Spaces:
Sleeping
Sleeping
Samuel Thomas
commited on
Commit
·
9c250a6
1
Parent(s):
0d26ceb
use cpu
Browse files
tools.py
CHANGED
|
@@ -2257,7 +2257,7 @@ def create_llm_pipeline():
|
|
| 2257 |
return pipeline(
|
| 2258 |
"text-generation",
|
| 2259 |
model=model_id,
|
| 2260 |
-
device_map="
|
| 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 |
|