tf32推理

2026-02-10 16:39:14 +08:00
parent fffc5a9956
commit ed637c972b
6 changed files with 16 additions and 2 deletions
--- a/scripts/trainer.py
+++ b/scripts/trainer.py
@@ -11,6 +11,9 @@ from unifolm_wma.utils.utils import instantiate_from_config
 from unifolm_wma.utils.train import get_trainer_callbacks, get_trainer_logger, get_trainer_strategy
 from unifolm_wma.utils.train import set_logger, init_workspace, load_checkpoints, get_num_parameters

+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+

 def get_parser(**parser_kwargs):
    parser = argparse.ArgumentParser(**parser_kwargs)