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 |
|