现在位置: 首页 > PyTorch 教程 > 正文

PyTorch torch.cuda.empty_cache 函数


Pytorch torch 参考手册 Pytorch torch 参考手册

torch.cuda.empty_cache 是 PyTorch 中用于清空 CUDA 缓存的函数。释放未使用的缓存内存,使其可以重新分配。

函数定义

torch.cuda.empty_cache()

使用示例

实例

import torch

# 清空 CUDA 缓存
if torch.cuda.is_available():
    # 创建一些张量
    x = torch.randn(1000, 1000).cuda()
    y = torch.randn(1000, 1000).cuda()

    # 清空缓存
    torch.cuda.empty_cache()
    print("CUDA cache cleared")

    # 检查内存
    print(f"Memory allocated: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
else:
    print("CUDA not available")

Pytorch torch 参考手册 Pytorch torch 参考手册