tf32推理

scripts/evaluation/base_model_inference.py

@@ -16,6 +16,9 @@ from collections import OrderedDict
 from unifolm_wma.models.samplers.ddim import DDIMSampler
 from unifolm_wma.utils.utils import instantiate_from_config
 
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
 
 def get_filelist(data_dir: str, postfixes: list[str]) -> list[str]:
     """

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
 

scripts/evaluation/world_model_interaction.py

@@ -18,6 +18,9 @@ from collections import OrderedDict
 from torch import nn
 from eval_utils import populate_queues, log_to_tensorboard
 from collections import deque
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
 from torch import Tensor
 from torch.utils.tensorboard import SummaryWriter
 from PIL import Image