Add a command line argument to enable backend:cudaMallocAsync

This commit is contained in:
comfyanonymous
2023-07-17 11:00:14 -04:00
parent 3a150bad15
commit 1679abd86d
3 changed files with 11 additions and 2 deletions

View File

@@ -204,7 +204,7 @@ print(f"Set vram state to: {vram_state.name}")
def get_torch_device_name(device):
if hasattr(device, 'type'):
if device.type == "cuda":
return "{} {}".format(device, torch.cuda.get_device_name(device))
return "{} {} : {}".format(device, torch.cuda.get_device_name(device), torch.cuda.get_allocator_backend())
else:
return "{}".format(device.type)
else: