Fix potential issue with multi devices.

This commit is contained in:
comfyanonymous 2024-08-20 00:31:04 -04:00
parent 5a69f84c3c
commit 1b3eee672c

View File

@ -979,7 +979,7 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True, ma
if torch.version.hip: if torch.version.hip:
return True return True
props = torch.cuda.get_device_properties("cuda") props = torch.cuda.get_device_properties(device)
if props.major >= 8: if props.major >= 8:
return True return True
@ -1035,7 +1035,7 @@ def should_use_bf16(device=None, model_params=0, prioritize_performance=True, ma
if is_intel_xpu(): if is_intel_xpu():
return True return True
props = torch.cuda.get_device_properties("cuda") props = torch.cuda.get_device_properties(device)
if props.major >= 8: if props.major >= 8:
return True return True