ROCm 6.0 seems to require get_device_name to be called before memory methods in order to return valid data

This commit is contained in:
doctorpangloss 2024-05-22 22:09:07 -07:00
parent bb159a219e
commit 35cf996b68
2 changed files with 5 additions and 3 deletions

View File

@ -132,13 +132,15 @@ def get_total_memory(dev=None, torch_total_too=False):
_, mem_total_cuda = torch.cuda.mem_get_info(dev)
mem_total_torch = mem_reserved
mem_total = mem_total_cuda
if torch_total_too:
return (mem_total, mem_total_torch)
return mem_total, mem_total_torch
else:
return mem_total
# we're required to call get_device_name early on to initialize the methods get_total_memory will call
if torch.cuda.is_available() and hasattr(torch.version, "hip") and torch.version.hip is not None:
logging.info(f"Detected HIP device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
total_vram = get_total_memory(get_torch_device()) / (1024 * 1024)
total_ram = psutil.virtual_memory().total / (1024 * 1024)
logging.info("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram))

View File

@ -162,7 +162,7 @@ def dependencies(force_nightly: bool = False) -> List[str]:
index_urls_selected = [stable for (stable, _) in index_urls]
_alternative_indices_selected = [stable for (stable, _) in _alternative_indices]
try:
# pip 23
# pip 23, 24
finder = PackageFinder.create(LinkCollector(session, SearchScope([], index_urls_selected, no_index=False)),
SelectionPreferences(allow_yanked=False, prefer_binary=False,
allow_all_prereleases=True))