tf32推理

2026-02-10 16:39:14 +08:00
parent fffc5a9956
commit ed637c972b
6 changed files with 16 additions and 2 deletions
--- a/scripts/evaluation/real_eval_server.py
+++ b/scripts/evaluation/real_eval_server.py
@@ -19,6 +19,9 @@ from fastapi.responses import JSONResponse
 from typing import Any, Dict, Optional, Tuple, List
 from datetime import datetime

+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
 from unifolm_wma.utils.utils import instantiate_from_config
 from unifolm_wma.models.samplers.ddim import DDIMSampler