监控显存使用

PyTorch提供了两个函数用于显存查询：

memory_allocated

torch.cuda.memory_allocated(device=None)

查询指定GPU中使用的显存大小（字节）

torch.cuda.max_memory_allocated(device=None)

返回给定设备的张量占用的最大GPU内存（字节）