Add env TORCH_AMD_CUDNN_ENABLED

This commit is contained in:
Christopher Anderson 2025-11-08 03:37:46 +11:00
parent a1a70362ca
commit 417dbcb32f

View File

@ -338,8 +338,11 @@ try:
if is_amd(): if is_amd():
arch = torch.cuda.get_device_properties(get_torch_device()).gcnArchName arch = torch.cuda.get_device_properties(get_torch_device()).gcnArchName
if not (any((a in arch) for a in AMD_RDNA2_AND_OLDER_ARCH)): if not (any((a in arch) for a in AMD_RDNA2_AND_OLDER_ARCH)):
torch.backends.cudnn.enabled = False # Seems to improve things a lot on AMD torch.backends.cudnn.enabled = os.environ.get("TORCH_AMD_CUDNN_ENABLED", "0").strip().lower() not in {
logging.info("Set: torch.backends.cudnn.enabled = False for better AMD performance.") "0", "off", "false", "disable", "disabled", "no"}
if not torch.backends.cudnn.enabled:
logging.info(
"ComfyUI has set torch.backends.cudnn.enabled to False for better AMD performance. Set environment var TORCH_AMD_CUDDNN_ENABLED=1 to enable it again.")
try: try:
rocm_version = tuple(map(int, str(torch.version.hip).split(".")[:2])) rocm_version = tuple(map(int, str(torch.version.hip).split(".")[:2]))