|
@@ -16,7 +16,7 @@ def enable_custom_logits_processors() -> bool:
|
|
|
major, minor = torch.cuda.get_device_capability()
|
|
major, minor = torch.cuda.get_device_capability()
|
|
|
# 正确计算Compute Capability
|
|
# 正确计算Compute Capability
|
|
|
compute_capability = f"{major}.{minor}"
|
|
compute_capability = f"{major}.{minor}"
|
|
|
- elif torch.npu.is_available():
|
|
|
|
|
|
|
+ elif hasattr(torch, 'npu') and torch.npu.is_available():
|
|
|
compute_capability = "8.0"
|
|
compute_capability = "8.0"
|
|
|
else:
|
|
else:
|
|
|
logger.info("CUDA not available, disabling custom_logits_processors")
|
|
logger.info("CUDA not available, disabling custom_logits_processors")
|