mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2026-04-26 09:42:27 +00:00
10 lines
240 B
Python
10 lines
240 B
Python
import torch
|
|
from colossalai.utils import get_current_device
|
|
|
|
|
|
def col_cuda_memory_capacity():
|
|
"""
|
|
Get cuda memory capacity of the current cuda.
|
|
"""
|
|
return torch.cuda.get_device_properties(get_current_device()).total_memory
|