Update model_management.py

This commit is contained in:
patientx 2024-08-16 00:53:38 +03:00 committed by GitHub
parent e58e586c29
commit 7bd86b0896
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -265,10 +265,24 @@ def get_torch_device_name(device):
return "CUDA {}: {}".format(device, torch.cuda.get_device_name(device))
try:
logging.info("Device: {}".format(get_torch_device_name(get_torch_device())))
except:
logging.warning("Could not pick default device.")
torch_device_name = get_torch_device_name(get_torch_device())
if "[ZLUDA]" in torch_device_name:
print("***----------------------ZLUDA--------------------------***")
print(" :: ZLUDA detected, disabling non-supported functions.")
if torch.backends.cudnn.enabled:
torch.backends.cudnn.enabled = False
print(" :: cuDNN , flash_sdp , mem_efficient_sdp disabled")
torch.backends.cuda.enable_flash_sdp(False)
torch.backends.cuda.enable_math_sdp(True)
torch.backends.cuda.enable_mem_efficient_sdp(False)
print("***-----------------------------------------------------***")
print("Device:", torch_device_name)
except:
print("Could not pick default device.")
current_loaded_models = []