fix(api): correct device for VRAM check
This commit is contained in:
parent
56a826228d
commit
cf8ca34845
|
@ -144,7 +144,7 @@ def run_gc(devices: List[DeviceParams] = []):
|
||||||
with torch.cuda.device(device.torch_str()):
|
with torch.cuda.device(device.torch_str()):
|
||||||
torch.cuda.empty_cache()
|
torch.cuda.empty_cache()
|
||||||
torch.cuda.ipc_collect()
|
torch.cuda.ipc_collect()
|
||||||
mem_free, mem_total = torch.cuda.mem_get_info(device=device.torch_str())
|
mem_free, mem_total = torch.cuda.mem_get_info()
|
||||||
logger.debug(
|
logger.debug(
|
||||||
"remaining CUDA VRAM usage: %s of %s",
|
"remaining CUDA VRAM usage: %s of %s",
|
||||||
(mem_total - mem_free),
|
(mem_total - mem_free),
|
||||||
|
|
Loading…
Reference in New Issue