feie9454/kronos-trader
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: add torch.cuda.empty_cache() during autoregressive inference

Browse Source 
Without releasing cached GPU memory, usage will keep growing during autoregressive prediction, leading to significant memory increase or OOM. Calling torch.cuda.empty_cache() prevents this accumulation.
This commit is contained in:
Pengxiao Song 2025-09-02 10:26:27 +08:00 committed by GitHub
parent 939986adb1
commit e027051b38
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
model/kronos.py
View File

@ -432,6 +432,8 @@ def auto_regressive_inference(tokenizer, model, x, x_stamp, y_stamp, max_context
x_token[0] = torch.cat([x_token[0], sample_pre], dim=1)
x_token[1] = torch.cat([x_token[1], sample_post], dim=1)
torch.cuda.empty_cache()
input_tokens = [t[:, -max_context:].contiguous() for t in x_token]
z = tokenizer.decode(input_tokens, half=True)
z = z.reshape(batch_size, sample_count, z.size(1), z.size(2))
@ -621,3 +623,4 @@ class KronosPredictor:
pred_dfs.append(pred_df)
return pred_dfs