add comments about empty_cache

Signed-off-by: daquexian <[email protected]>
BlinkDL · Jul 7, 2023 · 3d49dc2 · 3d49dc2
1 parent 1e26431
commit 3d49dc2
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/rwkv_pip_package/src/rwkv/model.py b/rwkv_pip_package/src/rwkv/model.py
@@ -703,6 +703,8 @@ def forward(self, tokens, state, full_output=False):
                     del pb[vw_key]
                     del pb[rw_key]
                     del pb[ow_key]
+                    # PyTorch's memory pool fails to reuse the freed memory.
+                    # So we empty the cache manually.
                     torch.cuda.empty_cache()
 
                 if prefetch: