[diffusion] fix: fix the bug of redundant memory usage on GPU-0 (#18221)
This commit is contained in:
@@ -124,6 +124,9 @@ class CudaPlatformBase(Platform):
|
|||||||
if empty_cache:
|
if empty_cache:
|
||||||
torch.cuda.empty_cache()
|
torch.cuda.empty_cache()
|
||||||
|
|
||||||
|
if torch.distributed.is_initialized():
|
||||||
|
device_id = torch.distributed.get_rank()
|
||||||
|
|
||||||
device_props = torch.cuda.get_device_properties(device_id)
|
device_props = torch.cuda.get_device_properties(device_id)
|
||||||
if device_props.is_integrated:
|
if device_props.is_integrated:
|
||||||
free_gpu_memory = psutil.virtual_memory().available
|
free_gpu_memory = psutil.virtual_memory().available
|
||||||
|
|||||||
Reference in New Issue
Block a user